Tensor operations

The last lesson left us pressed against the edge of two dimensions — a + b that added a vector to a whole batch, a hint of images stacked with height, width, and colour. This is the rung above the matrix. A scalar is a single number, a vector a 1-D list, a matrix a 2-D grid — and a tensor simply continues the ladder: an n-dimensional array.

Every modern ML framework (NumPy, PyTorch, JAX) is, at heart, a tensor library, and a neural network is nothing more exotic than a sequence of tensor operations — reshape, broadcast, multiply, reduce. Learn to read those four and the code stops looking like magic; miss one of them — almost always broadcasting — and you get the quietest, costliest bug in all of machine learning.

Rank, shape, axes

A tensor’s rank (or order) is how many indices you need to address an element: rank 0 = scalar, 1 = vector, 2 = matrix, 3+ = “tensor” proper. Its shape is the tuple of dimension sizes, and each dimension is an axis (numbered from 0).

A batch of 32 RGB images of 64×64 pixels is a rank-4 tensor of shape (32, 3, 64, 64); axis 0 is the batch, axis 1 the channels, axes 2–3 the height and width.

Two operations rearrange without changing the data:

• Reshape — reinterpret the same elements under a new shape (reshape, flatten, squeeze/unsqueeze). The element count is preserved; only the indexing changes.
• Reduction — collapse an axis with sum/mean/max along it: summing a (32, 10) tensor over axis 1 gives shape (32,).

Broadcasting — the rule everyone trips on

What happens when you add two tensors of different shapes? Broadcasting decides. The rule is:

• Align the shapes from the trailing axis.
• Two axes are compatible if they are equal or one of them is 1.
• A size-1 axis is virtually stretched to match.

A (3, 1) column plus a (1, 4) row therefore both stretch to (3, 4):

import numpy as np

col = np.array([[1], [2], [3]])      # (3, 1)
row = np.array([[10, 20, 30, 40]])   # (1, 4)
print("broadcast sum shape:", (col + row).shape)

A = np.arange(6).reshape(2, 3)
B = np.arange(6).reshape(3, 2)
print("einsum 'ij,jk->ik' == A@B:", np.array_equal(np.einsum("ij,jk->ik", A, B), A @ B))

X = np.ones((5, 2, 3)); Y = np.ones((5, 3, 4))
print("batched matmul shape:", (X @ Y).shape)   # contract last two dims, batch over the 5

broadcast sum shape: (3, 4)
einsum 'ij,jk->ik' == A@B: True
batched matmul shape: (5, 2, 4)

einsum and batched matmul

Einstein summation (einsum) is a compact notation for sums-of-products over named indices. Repeated indices are summed, and free indices stay.

• ij,jk->ik is matrix multiplication.
• ii-> is the trace.
• ij->ji is a transpose.

It’s the one notation that covers dot products, matmuls, batched matmuls, and contractions uniformly — and it makes the axes explicit, which kills broadcasting confusion.

Batched matmul is the everyday workhorse: @ on tensors of rank ≥ 3 multiplies the last two axes as matrices and broadcasts over the leading ones. That is exactly how attention computes scores Q @ Kᵀ over a (batch, heads, seq, dim) tensor. It is one operation, applied across every batch and head at once (see multi-head attention).

In one breath

• A tensor is an n-d array; rank = number of axes, shape = the size tuple, each axis numbered from 0 (a batch of images is (N, C, H, W)).
• Reshape rearranges the same elements; reductions (sum/mean/max along an axis) collapse a dimension.
• Broadcasting aligns shapes from the trailing axis — equal or one is 1 — and virtually stretches size-1 axes (so (3,1)+(1,4)→(3,4)). It is also the #1 silent bug when shapes mismatch unexpectedly.
• einsum is a uniform notation for dot/matmul/contraction (ij,jk->ik = matmul). Batched matmul (@ on rank ≥ 3) multiplies the last two axes and broadcasts the rest — exactly how attention runs Q @ Kᵀ over batch and heads.
• A neural net is just these operations composed — read the shapes and the mystery goes away.

Practice

A question to carry forward

Across these lessons we have learned to act on vectors — measure them, transform them, multiply and broadcast them. But there is an older, more basic question we have quietly stepped past, the very one linear algebra was invented to answer.

Multiplication runs a transform forward: given x, it hands you Ax. What about running it backward?

Suppose someone gives you the output and the transform and asks for the input: solve Ax = b for x. That is a system of linear equations — several constraints that must all hold at once — and the answer is not always a tidy “one solution.” Sometimes there is exactly one, sometimes none, sometimes infinitely many.

Here is the thread onward: how do you actually solve Ax = b by hand the way a computer does — row reduction to RREF? And what does the shape of the answer, one versus none versus infinitely many, quietly reveal about the matrix A itself?