PyTorch and NumPy idioms interviewers test: broadcasting and views, step by step

In this article

1. Before you start
2. Why interviewers ask this
3. PyTorch and NumPy idioms from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one step of the PyTorch interview questions and the machine learning (ML) coding round learning path.

Why interviewers ask this

A wrong shape rule costs a candidate the whole problem, because the code either crashes or silently returns an array of the wrong size. ML coding rounds therefore use short tensor tasks to check whether a candidate thinks in shapes. Candidates report prompts such as computing each point's distance to its closest center with broadcasting and argmin, converting a flat four-dimensional tensor from batch, channels, height, width (NCHW) order to batch, height, width, channels (NHWC) order, writing an L2 norm in NumPy, explaining why NumPy beats list multiplication, and implementing sparse matrix storage, addition and multiplication.

A weak answer reaches for nested loops and fixes shape errors by trial. A strong answer states the shape of every intermediate before typing. A Staff-level answer adds the memory cost of the intermediate and a cheaper alternative.

PyTorch and NumPy idioms from first principles

Vectorization

Vectorization is the practice of writing a computation as operations on whole arrays that the library runs in compiled code, so no Python loop touches individual elements. The loop still exists, but it runs in compiled code instead of Python (C on the central processing unit (CPU), and Compute Unified Device Architecture (CUDA) kernels on a graphics processing unit (GPU) in PyTorch), and that move is where most of the speed comes from.

Stride

A stride is the number of elements, or bytes, to skip in the underlying flat memory to move one step along one axis of an array. Transposing an array swaps its strides rather than moving its data, which is why a transpose is cheap. NumPy reports strides in bytes, while PyTorch .stride() reports elements.

In plain words

Picture a seating chart. You have a column of name tags down the left edge and a row of table numbers along the top. Pairing every tag with every table gives a grid. You do not print the tag again for every table. You read the same tag along its whole row. Broadcasting works the same way. The analogy breaks in one place. The repeated tags are only imagined, but the grid of results is real and takes memory. A mistake in which edge holds the tags shows up as a wrong shape, not as a visible picture.

In one precise sentence, these idioms let you describe a result by the shapes of the inputs and the axes you combine, and let the library do the looping.

Step 1: Vectorizing moves the loop into compiled code

Vectorizing replaces a Python loop with one array operation, and that is why NumPy beats list multiplication. The NumPy user guide, Broadcasting (page metadata last updated June 2026) describes vectorizing as arranging for looping to occur in C instead of Python. With lists, [x * y for x, y in zip(a, b)] runs interpreter code once per element, and a * b on two Python lists is a TypeError. With arrays, a * b hands both arrays to one compiled routine. An array also holds one number type in one block of memory, so the routine reads values directly instead of unpacking Python objects.

Input: two arrays of the same shape. Output: one array of products. You now know why the idiom exists; the next step explains how arrays of different shapes still combine.

Step 2: Broadcasting aligns shapes from the right

Broadcasting compares two shapes starting at the rightmost dimension and moving left, and two dimensions are compatible when they are equal or one of them is 1, as the NumPy broadcasting rules state. A missing leading dimension counts as size one. Shapes that fail raise a ValueError. The guide's example aligns an 8 by 1 by 6 by 1 array with a 7 by 1 by 5 array and gets an 8 by 7 by 6 by 5 result.

The stretching is conceptual: the guide says NumPy uses the original values without making needless copies. To pair every point with every center, insert a size-one axis with None: points[:, None, :] has shape (5, 1, 2) and centers[None, :, :] has shape (1, 3, 2) for five points and three centers (illustrative). Aligned from the right, they combine into (5, 3, 2).¹ The same indexing works in PyTorch. Scaled dot-product attention relies on the same rule; see implementing scaled dot-product attention from scratch.

Step 3: A view reuses memory, and strides show how

A view is an array object that reads the same memory as another array through a different shape or strides, so creating one copies nothing. A transpose is the standard example. For a (2, 3, 4, 5) NCHW array of 32-bit numbers, the strides are 60, 20, 5 and 1 elements.¹ Transposing to NHWC order, transpose(0, 2, 3, 1) in NumPy and permute(0, 2, 3, 1) in PyTorch, gives shape (2, 4, 5, 3) with strides 60, 5, 1 and 20, and the memory is still shared.¹ NumPy's a.strides reports bytes, as the NumPy ndarray reference describes, so for 4-byte numbers the NCHW array shows 240, 80, 20 and 4, while PyTorch .stride() reports the element counts above.¹

A contiguous array is one whose elements sit in memory in the order a row-by-row reading visits them. The transposed view is not contiguous, so code that treats memory as a flat NHWC buffer needs a copy: in NumPy np.ascontiguousarray, in PyTorch .contiguous() (see the PyTorch tensor views documentation). The copy has strides 60, 15, 3 and 1 and shares no memory.¹ The sliding-window loops of 2D convolution from scratch depend on this layout.

For a flat array of eight numbers holding a (1, 2, 2, 2) NCHW tensor, the NHWC order is 0, 4, 1, 5, 2, 6, 3, 7.¹ The index maps are ((n*C + c)*H + h)*W + w for NCHW and ((n*H + h)*W + w)*C + c for NHWC.

Step 4: Reducing along an axis removes that axis

A reduction collapses one axis of an array with an operation such as a sum or a minimum, and the axis you name disappears from the shape. In NumPy the argument is axis; in PyTorch it is dim. Summing the last axis of a (5, 3, 2) array leaves shape (5, 3).¹ A function called argmin returns the position of the smallest value along an axis instead of the value itself, so argmin along axis 1 of a (5, 3) array returns five center indices.¹

The same pattern gives an L2 norm: square, sum along the last axis, take the square root, as in np.sqrt((x * x).sum(axis=-1)). Library one-liners exist, np.linalg.norm(x, axis=-1) and torch.linalg.norm(x, dim=-1); an interviewer may accept or forbid them, so ask first. Integer-array indexing, called advanced indexing, picks entries by position: scores[rows, pos_ids] reads one value per row, and torch.gather does the same along a dimension.

It is the building block for negative-sampling prompts where the positive ids are known. Always say which axis disappears. A reduction over the wrong axis returns a plausible array of the wrong meaning, and shape checks will not catch it when two axes have the same length. Stable softmax is another reduction pattern, shown in a numerically stable softmax.

Step 5: Sparse storage keeps only the nonzero entries

Sparse storage keeps only the nonzero entries of a matrix. The coordinate list (COO) format is three parallel lists, one each for row index, column index and value. For a 3 by 4 matrix with entries 2 at (0, 1), 1 at (1, 0) and 3 at (1, 3), the lists are rows 0, 1, 1, columns 1, 0, 3 and values 2, 1, 3, which is 9 stored numbers against 12 dense cells (illustrative).¹ Compressed sparse row (CSR) replaces the row list with a pointer list, here 0, 1, 3, 3.¹

Addition merges entries by position and adds values that share a position. Adding a second illustrative matrix with three nonzeros gives five nonzeros.¹ Multiplying by a vector loops over stored entries and adds value * x[col] into y[row]; with x = (1, 2, 3, 4) the result is 4, 13, 0.¹ At an illustrative size of 100,000 by 100,000 with 1,000,000 nonzeros, dense 32-bit storage takes 40,000,000,000 bytes and COO with 64-bit indices takes 20,000,000 bytes.¹

1 · Vectorize2 · Broadcast3 · View4 · Reduce5 · SparseReplay

Step-by-step description

1. The loop moves into compiled code.
2. Aligned from the right, (5,1,2) and (1,3,2) give (5,3,2).
3. A transpose swaps strides, not data.
4. The named axis disappears.
5. Only nonzero entries are stored.

• Vectorizing moves the loop into compiled code.
• Broadcasting aligns shapes from the right.
• A transpose changes strides, not data.
• A reduction along an axis removes that axis.
• Sparse storage keeps only nonzero entries.

In short: vectorizing moves the loop into compiled code, broadcasting aligns shapes from the right, a transpose changes strides and not data, a reduction along an axis removes that axis, and sparse storage keeps only nonzero entries.

Worked example

The closest-center problem chains Steps 2 and 4. Take five points and three centers (illustrative integers, so you can redo it by hand).

import numpy as np

points = np.array([[0, 0], [1, 1], [4, 4], [5, 0], [2, 3]])
centers = np.array([[0, 1], [4, 3], [5, 1]])

diff = points[:, None, :] - centers[None, :, :]   # shape (5, 3, 2)
sq = (diff ** 2).sum(axis=-1)                      # shape (5, 3)
nearest = sq.argmin(axis=1)                        # shape (5,)
dist = np.sqrt(sq.min(axis=1))                     # shape (5,)

The PyTorch version swaps np.array for torch.tensor and axis for dim; the shapes are the same. Each point's squared distances to the three centers fill one row of sq; the first row is 1, 25, 26.¹

The last two rows are the senior-level point. For 1,000 points and 50 centers (illustrative), the difference tensor takes 400,000 bytes and the final distance matrix takes 200,000 bytes, so the intermediate is twice the size of the answer.¹

Text description of the diagram

• The difference tensor takes 400,000 bytes.
• The distance matrix takes 200,000 bytes.

Trade-offs and edge cases

Broadcasting trades memory for brevity. The NumPy guide itself warns that broadcasting can be a bad idea when it leads to inefficient use of memory. The difference tensor above grows with points times centers times dimensions. A cheaper form expands the square, the squared point norm plus the squared center norm minus twice their dot product, which needs only the points-by-centers matrix of 200,000 bytes in the example.¹ Rounding can make a squared distance slightly negative, so clamp at zero before the square root. For very large inputs, process the points in chunks.

Views trade safety for speed. Writing to a view changes the original array, because both read the same memory, and the transposed NHWC view above is not contiguous.¹ If you need independent data or a flat buffer, copy.

Text description of the diagram

• A transpose view shares memory with the original.
• A contiguous copy shares no memory.

Sparse storage trades simplicity for size. It wins only when most entries are zero, and merging entries by position costs more bookkeeping per stored entry than a dense add. Coordinate lists suit building a matrix; row-pointer formats suit multiplying by vectors. For code that builds whole models from these operations, see building generative pre-trained transformer (GPT) from scratch in PyTorch and multi-head attention without undefined.

How to answer it in an interview

The spoken structure: say the shape of each intermediate, name the axis you reduce, then give the memory cost and the cheaper alternative.

Common mistakes

• "Broadcasting copies the data." (wrong) The stretching is conceptual and NumPy avoids needless copies. (right) "Broadcasting reuses the values; the cost is the size of the result."
• "I compare shapes from the left." (wrong) The comparison starts at the rightmost dimension, so a (3,) array pairs with the last axis. (right) "I align from the right and treat missing dimensions as size one."
• "Transpose copies the array, so NHWC is a new buffer." (wrong) It is a view with swapped strides. (right) "The transpose is a view; I call a contiguous copy when I need flat NHWC order."
• "I sum over axis 0 because the code ran." (wrong) A square array hides the wrong axis. (right) "I say which axis disappears and check the output shape."

Follow-up questions

Why is numpy faster than list multiplication?

A list multiplication runs the interpreter once per element, while an array operation runs one compiled loop over a single block of one number type. Say "vectorization" and name the C loop; do not stop at "NumPy is written in C".

Implement sparse matrix storage, addition, and multiplication?

Store row, column and value lists (coordinate list), merge by position for addition, and for a vector product add value * x[col] into y[row] for each stored entry. State that the cost scales with the number of nonzeros, not the dense size.

How do you convert a flat NCHW tensor to NHWC?

Loop over n, c, h, w, read index ((n*C + c)*H + h)*W + w and write index ((n*H + h)*W + w)*C + c. In PyTorch, permute(0, 2, 3, 1) is the view and .contiguous() the copy.

What if the broadcast difference tensor does not fit in memory?

Expand the squared distance into norms and a matrix product, which avoids the points-by-centers-by-dimensions tensor, or loop over chunks of points and keep only each chunk's nearest center. Clamp small negative values at zero.

Walk through how a model goes from definition to GPU execution?

NumPy runs only on the CPU, while PyTorch in eager mode dispatches each operation on GPU tensors to a GPU kernel as you call it. Kernels launch asynchronously on a CUDA stream, and autograd records the graph for the backward pass. torch.compile instead captures a graph and fuses kernels. The jax library traces the function under jit and compiles it with Accelerated Linear Algebra (XLA) into fused device code. See GPU kernel questions in ML coding rounds for kernel detail and fully sharded data parallel (FSDP) in what FSDP is for several GPUs.

Sources

1. NumPy user guide: Broadcasting — NumPy project documentation, 28 Jun 2026
2. PyTorch documentation: Tensor Views — PyTorch documentation, 26 Oct 2026
3. NumPy reference: The N-dimensional array (ndarray) — NumPy project documentation, 26 Oct 2026