Tensors as nested lists

A list with a contract

A vector is a list of numbers. A matrix is a list of lists. A batch of colour images is a list of lists of lists of lists. At that depth, the nesting stops being a convenient picture and becomes the data structure. A tensor is the general name for that structure. The numbers are the easy part. The shape is the contract, and it says which axis means what.

The rank is the number of axes. The shape is the length along each one. A batch of 32 RGB images at 224 by 224 pixels is (32, 3, 224, 224) in PyTorch. The same batch is (32, 224, 224, 3) in TensorFlow. The pixels are identical; the promise about axis order is not. Code that ignores the promise still runs. That silence is what makes the bugs expensive.

A scalar has rank zero, a vector rank one, and a matrix rank two. The names matter less than the count. What matters is which axis holds the batch, which axis holds the channels, and whether the last axis is a feature or a class. Write those three facts down before you write a layer.

Axis order is a decision

PyTorch puts channels before height and width. A convolution reads a whole channel at a time, and contiguous channels make that read cheap. TensorFlow and most image libraries put channels last. That is how a camera and a PNG store them. Neither order is wrong. Mixing them is.

The batch axis comes first almost everywhere. One example and a hundred examples then travel through the same code. That convenience has a trap. Indexing tensor[0] removes the batch axis, so a layer that expected a batch now sees a different rank. Keep the axis with tensor[0:1] when you mean a batch of one.

Memory layout is the other half of the story. A contiguous tensor stores its values in row-major order, so the last axis varies fastest. Its strides say how many elements to skip when you move one step along each axis. A transpose usually returns a view with different strides, not a copy. A reshape after a transpose may have to copy, because the old strides cannot describe the new shape. The order of the two operations can then change both the speed and the result.

Shape, strides and a batched product

import numpy as np

x = np.arange(24).reshape(2, 3, 4)      # batch, rows, columns
print(x.shape, x.strides)

y = x.transpose(0, 2, 1)                # batch, columns, rows
print(y.shape)                          # (2, 4, 3)

z = np.einsum("bij,bjk->bik", x, y)
print(z.shape)                          # (2, 3, 3)

Reshape, transpose, squeeze

A reshape reinterprets the same flat buffer under a new shape. It reads the buffer in row-major order, so it does not move numbers and costs nothing. It will also produce a nonsensical arrangement if the order you had in mind was different. A transpose does move numbers, or at least changes how they are addressed. It is the operation that turns channels-first into channels-last.

Squeeze and unsqueeze are reshapes that add or remove one axis. They are the usual fix for a missing batch axis. They are also the usual cause of one. A model trained on (N, C) receives (C,) and treats the channels as a batch. Print the shape at every boundary, and this class of bug disappears.

Indexing follows the same logic. A slice keeps the axis, an integer index removes it, and a boolean mask flattens the result. tensor[:, 0] selects the first column of every row and drops the column axis. tensor[:, 0:1] keeps it. The difference is one character and one rank. The layer downstream will notice.

Broadcasting is a shortcut with rules

When two shapes meet, NumPy and every framework compare their axes from the right. Two axes are compatible if they are equal, or if one of them is 1. A size-1 axis is stretched to match the other side. A (2, 3) matrix plus a (3,) vector adds the vector to every row. A (2, 1) column plus a (1, 3) row produces a (2, 3) outer sum.

Broadcasting is why a bias vector needs no loop. It is also why a mistake can pass unnoticed. A (32, 3, 224, 224) batch minus a (3,) mean subtracts the mean from the last axis, which is width, not channels. The operation succeeds and the image turns into noise. Check which axis you meant, then reshape the statistic to (1, 3, 1, 1) so the intent is explicit.

A useful habit is to write the intended shapes beside the expression before running it. If the comment and the error disagree, either the comment is wrong or the code is. Either way, the bug is localised. Broadcasting is a convenience. It is not a substitute for knowing which axis you are on.