Tensors, Shapes, and Broadcasting

Different modalities reuse a small vocabulary of semantic dimensions: batch, sequence or time, channel or feature, spatial axes, head or expert. The vocabulary exists in every framework. What does not exist is any mechanism that holds a name to its axis. The array stores sizes and offsets. The meaning lives in your head, or in a comment.

Three researchers proposed writing the names into the mathematics itself, in a 2023 paper in Transactions on Machine Learning Research: “We propose a notation for tensors with named axes, which relieves the author, reader, and future implementers of machine learning models from the burden of keeping track of the order of axes and the purpose of each.”

Position is what the machine tracks. Meaning is what you have to track. The rest of this lesson is about what happens when you stop.

In applied ML, a tensor is a multidimensional array whose axes carry meaning. A shape such as 32×128×768 might mean batch, token, and hidden dimension. The numbers alone are ambiguous. Swapping the token and hidden axes preserves the element count and changes the computation completely.

Tensor literacy is therefore semantic bookkeeping. Every operation should state which axes it preserves, reorders, expands, reduces, or contracts — and, as the strides below show, which of those changes touch memory at all.

A contraction multiplies entries and sums over one or more matched axes. Matrix multiplication contracts the shared inner dimension. Einstein notation writes repeated indices as summed, allowing compact expressions such as yᵢ = Aᵢⱼxⱼ. The notation highlights which indices remain in the output. Attention scores, convolutions, and batched linear maps are all structured contractions.

Batched linear maps are where that index logic became a language feature, with a written rule about which axes are contracted and which merely come along. PEP 465 added Python's @ operator, and its specification states the rule: “For inputs with more than 2 dimensions, we treat the last two dimensions as being the dimensions of the matrices to multiply, and ‘broadcast’ across the other dimensions.” Its worked example is small enough to check by hand — “arr(10, 2, 3) @ arr(10, 3, 4) performs 10 separate matrix multiplies … returns the 10 resulting matrices together in an array with shape (10, 2, 4)”.

NumPy's reference page for numpy.matmul repeats it — “If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly” — and then shows what the older function does with identical operands. On a = np.ones([9,5,7,4]) and c = np.ones([9,5,4,3]), np.dot(a, c).shape is (9, 5, 7, 9, 5, 3) while np.matmul(a, c).shape is (9, 5, 7, 3). Two functions, the same two arrays, two perfectly legal output shapes. Only one of them is the batched contraction you meant.

Python 3.5 shipped the operator in 2015 and recorded the addition plainly: “PEP 465 adds the @ infix operator for matrix multiplication”, and “NumPy 1.10 has support for the new operator”. Libraries provide the optimized kernels. The index logic stays yours.

These four operations can share output dimensions while meaning very different things, and only one of them has a published contract you can cite when a shape surprises you. NumPy's user guide gives it: “When operating on two arrays, NumPy compares their shapes element-wise. It starts with the trailing (i.e. rightmost) dimension and works its way left.” Two axes are compatible when “they are equal, or” when “one of them is 1”. Otherwise the operation stops with the exact message “ValueError: operands could not be broadcast together”. Where it succeeds, it proceeds “without making needless copies of data”.

That right-to-left rule is not a NumPy convenience. The Python array API standard, in its 2025.12 text, defines broadcasting as “the automatic (implicit) expansion of array dimensions to be of equal sizes without copying array data” and writes out the rank-promotion algorithm step by step (“If d1 == 1, then set the i-th element of shape to d2” …). The ONNX operator specification adopts it wholesale — “Multidirectional broadcasting is the same as Numpy's broadcasting” — and tabulates worked cases such as shape(A) = (1, 4, 5), shape(B) = (2, 3, 1, 1) ==> shape(result) = (2, 3, 4, 5).

Three independent bodies, one rule. And the rule aligns axes by position rather than by meaning. Reshape, transpose and reduce have no comparable published contract about what your axes mean. They have only your intent, which is why the columns below have to be read as semantics and not as shapes.

The Transformer's own numbers make the trace concrete. The 2017 paper that introduced it fixes the count: “In this work we employ h = 8 parallel attention layers, or heads.” At d_model = 512 that fixes d_k = d_v = d_model/h = 64, and the operation itself is Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))V. So the trace runs over batch, 8 heads, tokens, and a 64-wide feature axis — not over an abstract d.

This failure has been written up with real shapes rather than a hypothetical one. Alexander Rush opens his account bluntly: “Despite its ubiquity in deep learning, Tensor is broken. It forces bad habits such as exposing private dimensions, broadcasting based on absolute position, and keeping type information in documentation.”

The demonstration is a rotate() implemented as ims.transpose(1, 2). On the batch it was written against, ims.shape is torch.Size([6, 96, 96, 3]) and the function does what its name says. Handed a single image, ims[0], the same call returns torch.Size([96, 3, 96]). The three-wide channel axis is now in the middle, sitting where a spatial axis used to be. Nothing raised.

Write this trace before implementing an unfamiliar tensor expression. The np.dot and np.matmul pair above is the reason it is worth the minute it costs: identical operands, (9, 5, 7, 9, 5, 3) against (9, 5, 7, 3), both legal, one of them silently the wrong computation. Steps one and two are what separate those two outcomes. They are cheaper than debugging a model that trains to a mediocre loss for a reason no exception ever reported.