Lecture 03: The Shapes of Data

The idea in one sentence

An array’s shape is the list of its axes, and nearly every operation in a numerical library is a rule about which axis disappears.

From one number to a stack of blocks

Most of the errors you will ever meet in numerical code are shape errors. The library stops, says two numbers failed to line up, and leaves you to work out which axis you actually meant. Learning to read shapes is how you stop losing afternoons to that message.

Lecture 2 left you with a matrix as a machine. In code that machine is an array, and the first thing an array knows about itself is its shape. What the numbers mean comes later, and it comes from you.

Start from nothing and add one axis at a time. A single number has no axes. Line numbers up and you have one axis. Stack those lines and you have two. Stack the stacks and you have three, and there is no reason to stop.

Fig. 1

The ladder. Every cell holds one number, and each rung adds one axis. The shape is just the lengths of the axes read from the outside in. The last rung is a real one: 32 color images of 224 by 224 pixels, the standard meal a vision model eats.

Two words carry the whole picture. The shape is the tuple of axis lengths. The number of axes has its own name, ndim, and it is the length of that tuple. Notice that shape says everything and means nothing: (32, 3, 224, 224) does not know that 3 is color channels. You know that.

python

import numpy as np

s = np.array(7.0)                       # 0D
v = np.array([1., 2., 3., 4.])          # 1D
M = np.array([[1., 2., 3.], [4., 5., 6.]])   # 2D
T = np.arange(24).reshape(2, 3, 4)      # 3D
B = np.zeros((32, 3, 224, 224))         # 4D

for a in (s, v, M, T, B):
    print(a.ndim, a.shape, a.size)

# 0 ()                 1
# 1 (4,)               4
# 2 (2, 3)             6
# 3 (2, 3, 4)          24
# 4 (32, 3, 224, 224)  4816896

python

import torch

s = torch.tensor(7.0)                        # 0D
v = torch.tensor([1., 2., 3., 4.])           # 1D
M = torch.tensor([[1., 2., 3.], [4., 5., 6.]])   # 2D
T = torch.arange(24).reshape(2, 3, 4)        # 3D
B = torch.zeros(32, 3, 224, 224)             # 4D

for a in (s, v, M, T, B):
    print(a.dim(), tuple(a.shape), a.numel()) # dim not ndim, numel not size

# 0 ()                 1
# 1 (4,)               4
# 2 (2, 3)             6
# 3 (2, 3, 4)          24
# 4 (32, 3, 224, 224)  4816896

python

import jax.numpy as jnp

s = jnp.array(7.0)                       # 0D
v = jnp.array([1., 2., 3., 4.])          # 1D
M = jnp.array([[1., 2., 3.], [4., 5., 6.]])   # 2D
T = jnp.arange(24).reshape(2, 3, 4)      # 3D
B = jnp.zeros((32, 3, 224, 224))         # 4D

for a in (s, v, M, T, B):
    print(a.ndim, a.shape, a.size)       # same names as NumPy, by design

# 0 ()                 1
# 1 (4,)               4
# 2 (2, 3)             6
# 3 (2, 3, 4)          24
# 4 (32, 3, 224, 224)  4816896

python

import tensorflow as tf

s = tf.constant(7.0)                        # 0D
v = tf.constant([1., 2., 3., 4.])           # 1D
M = tf.constant([[1., 2., 3.], [4., 5., 6.]])   # 2D
T = tf.reshape(tf.range(24), (2, 3, 4))     # 3D
B = tf.zeros((32, 3, 224, 224))             # 4D

for a in (s, v, M, T, B):
    print(len(a.shape), tuple(a.shape), tf.size(a).numpy())

# 0 ()                 1
# 1 (4,)               4
# 2 (2, 3)             6
# 3 (2, 3, 4)          24
# 4 (32, 3, 224, 224)  4816896

The four tabs print the same numbers and use four different words for the same idea. PyTorch says dim() where NumPy says ndim, and numel() where NumPy says size. TensorFlow keeps tf.rank for the count and reads lengths off shape. These are not cosmetic. Half the time you spend reading someone else’s model goes into translating this vocabulary, so learn all of it once and stop paying.

The axis you name is the axis that disappears

Names out of the way. The operation you will write most often takes a block of numbers and hands back fewer of them, and it always works the same way.

Take the 3D block and look at it honestly. np.arange(24).reshape(2, 3, 4) holds the numbers 0 to 23, and the shape says how they are nested: two groups, each holding three rows, each row holding four numbers.

Fig. 2

One array, drawn flat. Axis 2 runs along a row and axis 1 runs down the rows, both of them visible at a glance. Axis 0 is the blue one, the axis a flat picture has to fake: it steps from the left grid to the right one. The entry at position $(b, i, j)$ is $12b + 4i + j$ .

Now reduce it. Every reduction takes a bundle of numbers and returns one, and the axis you name in the call is the axis whose numbers get bundled. That axis is gone from the answer:

(2, 3, 4) \xrightarrow{\;\text{axis }0\;} (3, 4), \qquad (2, 3, 4) \xrightarrow{\;\text{axis }1\;} (2, 4), \qquad (2, 3, 4) \xrightarrow{\;\text{axis }2\;} (2, 3).

Fig. 3

One array, three reductions. Cyan is what goes in, blue is the reduction, green is what comes out. Naming an axis is asking for it to be collapsed, and the result keeps every axis you did not name, in the order they already had.

Read one of them slowly. Summing along axis 2 bundles each row of four numbers into one, so the first row $0+1+2+3 = 6$ and the second $4+5+6+7 = 22$ . The whole answer is

\texttt{T.sum(axis=2)} = \begin{bmatrix} 6 & 22 & 38 \\ 54 & 70 & 86 \end{bmatrix},

with shape $(2, 3)$ , the original shape with the 4 removed. Summing along axis 0 instead bundles the two grids entry by entry, so the top-left result is $0 + 12 = 12$ , and the shape is $(3, 4)$ .

Try it. Name an axis and watch it go

(2, 3, 4) → (2, 3)

The yellow cells are the ones being bundled; the green cell is the single value they become. Count the yellow ones: there are always exactly as many as the length of the axis you named.

Two consequences follow, and both bite people in practice. Naming several axes collapses all of them: T.sum(axis=(0, 2)) leaves shape $(3,)$ and gives $(60, 92, 124)$ . Naming none collapses everything down to a 0D array, and T.sum() is 276.

The other consequence is what keepdims is for. If you want to subtract each row’s mean from that row, the mean has shape $(2, 4)$ while the block has shape $(2, 3, 4)$ , and those do not line up. Ask for keepdims=True and the collapsed axis stays as a length-1 stub, shape $(2, 1, 4)$ , which lines up with $(2, 3, 4)$ and gets stretched across the missing rows. A length-1 axis is a promise that the same value applies all the way along, and that promise is what makes broadcasting work.

Fig. 4

What a length-1 axis buys you. The orange row is the only one that exists in memory, and the dashed orange rows above and below it are that same row read again. Subtracting a mean from a whole block therefore costs no extra storage.

python

import numpy as np

T = np.arange(24).reshape(2, 3, 4)

for ax in (0, 1, 2):
    print(ax, T.sum(axis=ax).shape, T.sum(axis=ax).ravel())

# 0 (3, 4) [12 14 16 18 20 22 24 26 28 30 32 34]
# 1 (2, 4) [12 15 18 21 48 51 54 57]
# 2 (2, 3) [ 6 22 38 54 70 86]

print(T.sum(axis=(0, 2)))    # [ 60  92 124], shape (3,)
print(T.sum())               # 276, shape ()
print(T.min(axis=2))         # [[ 0  4  8] [12 16 20]]

m = T.mean(axis=1, keepdims=True)
print(m.shape, (T - m).shape)     # (2, 1, 4) (2, 3, 4), broadcast back

python

import torch

T = torch.arange(24).reshape(2, 3, 4)

for ax in (0, 1, 2):
    r = T.sum(dim=ax)
    print(ax, tuple(r.shape), r.flatten().tolist())

# 0 (3, 4) [12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34]
# 1 (2, 4) [12, 15, 18, 21, 48, 51, 54, 57]
# 2 (2, 3) [6, 22, 38, 54, 70, 86]

print(T.sum(dim=(0, 2)))     # [ 60,  92, 124], shape (3,)
print(T.sum())               # 276, shape ()
print(torch.amin(T, dim=2))  # [[ 0,  4,  8], [12, 16, 20]]
# careful: torch.min(T, dim=2) hands back values AND indices

m = T.float().mean(dim=1, keepdim=True)   # keepdim, with no s
print(tuple(m.shape), tuple((T - m).shape))    # (2, 1, 4) (2, 3, 4)

python

import jax.numpy as jnp

T = jnp.arange(24).reshape(2, 3, 4)

for ax in (0, 1, 2):
    print(ax, T.sum(axis=ax).shape, T.sum(axis=ax).ravel())

# 0 (3, 4) [12 14 16 18 20 22 24 26 28 30 32 34]
# 1 (2, 4) [12 15 18 21 48 51 54 57]
# 2 (2, 3) [ 6 22 38 54 70 86]

print(T.sum(axis=(0, 2)))    # [ 60  92 124], shape (3,)
print(T.sum())               # 276, shape ()
print(T.min(axis=2))         # [[ 0  4  8] [12 16 20]]

m = T.mean(axis=1, keepdims=True)
print(m.shape, (T - m).shape)     # (2, 1, 4) (2, 3, 4), broadcast back

python

import tensorflow as tf

T = tf.reshape(tf.range(24), (2, 3, 4))

for ax in (0, 1, 2):
    r = tf.reduce_sum(T, axis=ax)
    print(ax, tuple(r.shape), tf.reshape(r, [-1]).numpy())

# 0 (3, 4) [12 14 16 18 20 22 24 26 28 30 32 34]
# 1 (2, 4) [12 15 18 21 48 51 54 57]
# 2 (2, 3) [ 6 22 38 54 70 86]

print(tf.reduce_sum(T, axis=[0, 2]))  # [ 60  92 124], shape (3,)
print(tf.reduce_sum(T))               # 276, shape ()
print(tf.reduce_min(T, axis=2))       # [[ 0  4  8] [12 16 20]]

m = tf.reduce_mean(tf.cast(T, tf.float32), axis=1, keepdims=True)
print(tuple(m.shape))                 # (2, 1, 4)

Same numbers, four spellings. NumPy and JAX say axis, PyTorch says dim, TensorFlow puts reduce_ on the front of every one. The one that has actually cost people hours is torch.min(T, dim=2), which returns the values together with the positions they came from, so unpacking it as a single array quietly fails. torch.amin is the one that behaves like np.min.

Two different things called rank

One word in this vocabulary is booby-trapped, and it is a word you will meet constantly.

Frameworks call an array a tensor and call its number of axes the tensor’s rank. A scalar is rank 0, a vector is rank 1, a matrix is rank 2, a block is rank 3. That is ndim wearing a different hat, and tf.rank returns exactly that number.

Linear algebra also has a word rank, and it means something else entirely: the number of genuinely independent columns in a matrix. Part III of this course spends several lectures on it, because that number is what decides whether a system has one solution, none, or infinitely many.

Three products that look alike

NumPy ships three functions that all multiply and add. Telling them apart is a question about axes.

Start where they agree. For two 1D arrays, inner and dot both do the schoolbook thing: multiply matching entries and add:

(1, 2, 3, 4) \cdot (5, 6, 7, 8) = 5 + 12 + 21 + 32 = 70 .

outer does the opposite. It adds nothing and multiplies everything, giving all sixteen products $a_i b_j$ arranged in a 4 by 4 array. Nothing collapses, so both axes survive.

Fig. 5

The zoo, drawn as shapes: cyan and orange go in, green comes out. The only difference between the first two rows is which way the orange operand is turned, and turning it is exactly a transpose.

Now the place they part company, which is worth memorizing because the error message will not tell you. Take two arrays of the same shape $(2, 4)$ :

A = \begin{bmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \end{bmatrix}, \qquad B = \begin{bmatrix} 11 & 12 & 13 & 14 \\ 15 & 16 & 17 & 18 \end{bmatrix}.

np.inner(A, B) runs and gives a $2 \times 2$ answer. np.dot(A, B) raises. The rule behind both is one sentence each. inner contracts the last axis of the first argument with the last axis of the second, and both are length 4, so it works out. dot contracts the last axis of the first with the second-to-last of the second, and $B$ ‘s second-to-last axis has length 2, which does not match 4.

So for 2D arrays,

\texttt{np.inner}(A, B) = AB^{\mathsf{T}}, \qquad \texttt{np.dot}(A, B) = AB .

Both calls run when $B$ is square, and then they agree whenever $B$ equals its own transpose. The $B$ above is not square, so only one of them runs at all. Here is the answer inner gives, which you can check by hand: the top-left entry is row 1 of $A$ against row 1 of $B$ , or $1(11) + 2(12) + 3(13) + 4(14) = 130$ .

\texttt{np.inner}(A, B) = \begin{bmatrix} 130 & 170 \\ 330 & 434 \end{bmatrix} .

python

import numpy as np

a = np.array([1., 2., 3., 4.])
b = np.array([5., 6., 7., 8.])
print(np.inner(a, b), np.dot(a, b))   # 70.0 70.0   agree in 1D
print(np.outer(a, b).shape)           # (4, 4)      nothing collapses

A = np.array([[1., 2., 3., 4.], [5., 6., 7., 8.]])
B = np.array([[11., 12., 13., 14.], [15., 16., 17., 18.]])
print(np.inner(A, B))          # [[130. 170.] [330. 434.]]
print(np.allclose(np.inner(A, B), A @ B.T))   # True: inner is A @ B.T

# np.dot(A, B) raises: shapes (2,4) and (2,4) not aligned
C = B.T                        # (4, 2), and now the 4s line up
print(np.dot(A, C))            # [[130. 170.] [330. 434.]]

python

import torch

a = torch.tensor([1., 2., 3., 4.])
b = torch.tensor([5., 6., 7., 8.])
print(torch.inner(a, b), torch.dot(a, b))   # 70.0 70.0   agree in 1D
print(torch.outer(a, b).shape)              # (4, 4)      nothing collapses

A = torch.tensor([[1., 2., 3., 4.], [5., 6., 7., 8.]])
B = torch.tensor([[11., 12., 13., 14.], [15., 16., 17., 18.]])
print(torch.inner(A, B))       # [[130., 170.], [330., 434.]]
print(torch.allclose(torch.inner(A, B), A @ B.T))   # True

# A @ B is not a legal matmul: (2,4) against (2,4)
C = B.T                        # (4, 2)
print(A @ C)                   # [[130., 330.], [170., 434.]]

python

import jax.numpy as jnp

a = jnp.array([1., 2., 3., 4.])
b = jnp.array([5., 6., 7., 8.])
print(jnp.inner(a, b), jnp.dot(a, b))   # 70.0 70.0   agree in 1D
print(jnp.outer(a, b).shape)            # (4, 4)      nothing collapses

A = jnp.array([[1., 2., 3., 4.], [5., 6., 7., 8.]])
B = jnp.array([[11., 12., 13., 14.], [15., 16., 17., 18.]])
print(jnp.inner(A, B))         # [[130. 170.] [330. 434.]]
print(jnp.allclose(jnp.inner(A, B), A @ B.T))   # True

# jnp.dot(A, B) raises for the same reason NumPy does
C = B.T                        # (4, 2)
print(jnp.dot(A, C))           # [[130. 330.] [170. 434.]]

python

import tensorflow as tf

a = tf.constant([1., 2., 3., 4.])
b = tf.constant([5., 6., 7., 8.])
print(tf.einsum('i,i->', a, b))       # 70.0        there is no tf.inner
print(tf.einsum('i,j->ij', a, b).shape)   # (4, 4)   and no tf.outer

A = tf.constant([[1., 2., 3., 4.], [5., 6., 7., 8.]])
B = tf.constant([[11., 12., 13., 14.], [15., 16., 17., 18.]])
print(tf.einsum('ik,jk->ij', A, B))   # [[130. 170.] [330. 434.]]

# A @ B is not a legal matmul: (2,4) against (2,4)
print(A @ tf.transpose(B))      # [[130. 170.] [330. 434.]], the 4s line up

einsum, the one language all four speak

Look at what the TensorFlow tab had to do. It has no inner and no outer, so it wrote both with einsum, and the strings it used are readable once you know the convention: name every axis with a letter, name the axes you want in the output after the arrow, and any letter that does not survive the arrow gets summed over.

That is the whole rule.

'i,i->' names one axis in each input, keeps nothing, so i is summed. That is the dot product.
'i,j->ij' names two different axes and keeps both, so nothing is summed. That is outer.
'ik,jk->ij' shares k between the inputs and drops it. That is inner, which is $AB^{\mathsf{T}}$ .
'ik,kj->ij' shares k in the other position. That is ordinary matrix multiplication.
'bik,kj->bij' adds a batch letter that both survives and is not shared. That is one matrix applied to a whole stack.

Fig. 6

Two strings, drawn. The yellow edges are the $k$ axis, the one both inputs share and the output does not have, so those are the edges that meet and vanish. Give every axis its letter and the rule reads itself.

The same five strings work unchanged in NumPy, PyTorch, JAX, and TensorFlow. When you are reading a model and cannot tell what a line of index gymnastics does, rewrite it as an einsum and the answer is in the string.

Five batches of ten images

Here is where all of this was going.

Before the shapes, look at what is actually inside them. Here is a hand-drawn 7 on an 8 by 8 grid, twice. On the left it is a picture. On the right it is the same 64 numbers, written out.

Fig. 7

The same object, drawn twice. The more heavily a cell is marked, the larger its number: 0 is untouched and 240 is the heaviest stroke here. The yellow cell is the same cell in both grids, and it holds 210. Nothing is lost going from left to right, which is the whole reason a computer can see.

An image is not kept as a picture. It is kept as the grid on the right, and the picture is what happens when you shade each number. Every operation in this lecture has been running on objects like this one all along.

Now stack them. Flatten each 8 by 8 grid into one row of 64 numbers, lay three rows on top of each other, and you have shape $(3, 64)$ : three images, 64 numbers each.

Fig. 8

Three pictures, three rows. Each row is its picture read left to right, top row of cells first, so the heavy run near the start of the top row is the bar across the 7. Blue marks the batch axis, the one that counts images instead of pixels. Read the shape from the outside in: 3 first, then 64.

The real numbers are only bigger. Swap 8 by 8 for 28 by 28 and a row of 64 becomes a row of 784. Swap three images for ten, then stack five of those, and the shape reads $(5, 10, 784)$ .

You have a small pile of images. Each one is 28 by 28 pixels, flattened into a row of 784 numbers. Ten of them make a batch, and you have five batches waiting, so the whole pile is one array of shape $(5, 10, 784)$ . The layer you want to push them through has 10000 output units, so its weights are a matrix of shape $(784, 10000)$ .

Write X @ W and every image goes through in one call:

(5, 10, 784) \; @ \; (784, 10000) \;\longrightarrow\; (5, 10, 10000) .

Fig. 9

One forward pass, drawn as shapes. Cyan is the images going in, orange is the one weight matrix they all share, green is what comes out. The 784 across the bottom of the left stack meets the 784 down the side of the weight matrix, those two cancel, and every other axis is carried along untouched.

The weight matrix has no batch axis, and it does not need one. The same 784 by 10000 matrix applies to every image in every batch, and the frameworks arrange that for free by matching axes from the right. This is broadcasting again, one level up: an axis that is absent behaves like an axis that is shared.

The two batch axes come out untouched, in the order they went in. The rule for @ on stacked arrays is that the last two axes do a matrix multiplication and everything in front of them is carried along. So $(5, 10, 784)$ against $(784, 10000)$ gives $(5, 10, 10000)$ , and the answer to “which image is row 7 of batch 3” is still Y[3, 7].

And the work is real. That single call does $5 \times 10 \times 784 \times 10000$ multiply-adds, which is 392 million of them, and the reason it takes a fraction of a second is that no Python loop ever runs. The shape told the library exactly which numbers to hand to which core, and the library did the rest.

python

import numpy as np

rng = np.random.default_rng(0)
X = rng.standard_normal((5, 10, 784)).astype(np.float32)  # 5 batches of 10
W = rng.standard_normal((784, 10000)).astype(np.float32)  # shared matrix

Y = X @ W
print(X.shape, W.shape, Y.shape)
# (5, 10, 784) (784, 10000) (5, 10, 10000)

# it really is one dot product per output number
print(np.allclose(Y[3, 7, 42], X[3, 7] @ W[:, 42], atol=1e-2))   # True
print(np.allclose(Y, np.einsum('bij,jk->bik', X, W), atol=1e-2)) # True
print(np.allclose(Y, np.inner(X, W.T), atol=1e-2))           # the long way

python

import torch

torch.manual_seed(0)
X = torch.randn(5, 10, 784)      # 5 batches, 10 images
W = torch.randn(784, 10000)      # one shared matrix

Y = X @ W
print(X.shape, W.shape, Y.shape)
# (5, 10, 784) (784, 10000) (5, 10, 10000)

print(torch.allclose(Y[3, 7, 42], X[3, 7] @ W[:, 42], atol=1e-2))  # True
e = torch.einsum('bij,jk->bik', X, W)
print(torch.allclose(Y, e, atol=1e-2))                             # True

python

import jax
import jax.numpy as jnp

k1, k2 = jax.random.split(jax.random.key(0))
X = jax.random.normal(k1, (5, 10, 784))    # 5 batches, 10 images
W = jax.random.normal(k2, (784, 10000))    # one shared matrix

Y = X @ W
print(X.shape, W.shape, Y.shape)
# (5, 10, 784) (784, 10000) (5, 10, 10000)

print(jnp.allclose(Y[3, 7, 42], X[3, 7] @ W[:, 42], atol=1e-2))    # True
e = jnp.einsum('bij,jk->bik', X, W)
print(jnp.allclose(Y, e, atol=1e-2))                               # True

python

import tensorflow as tf

tf.random.set_seed(0)
X = tf.random.normal((5, 10, 784))    # 5 batches, 10 images
W = tf.random.normal((784, 10000))    # one shared matrix

Y = X @ W
print(X.shape, W.shape, Y.shape)
# (5, 10, 784) (784, 10000) (5, 10, 10000)

one = tf.tensordot(X[3, 7], W[:, 42], 1)
print(tf.abs(Y[3, 7, 42] - one) < 1e-2)                    # True
e = tf.einsum('bij,jk->bik', X, W)
print(tf.reduce_all(tf.abs(Y - e) < 1e-2))                 # True

Where this is going

A shape settles whether an operation is legal. Whether it is worth doing is a separate question, and the shape stays quiet about that one. The weight matrix above turns 784 numbers into 10000, and nothing in its shape says whether those 10000 outputs carry 10000 independent pieces of information or secretly repeat the same 784. That question lives in the columns, and answering it is the job of the rest of this course. It starts with the one procedure that strips every matrix down to its skeleton: elimination, in the next lecture.

Check yourself

Without running anything, give the shape of np.arange(60).reshape(3, 4, 5).sum(axis=1), and then of the same array’s sum(axis=(0, 2)).
For the block $T$ in this lecture, compute T.max(axis=1) by hand and check it against the explorer above.
Two arrays have shapes $(2, 4)$ and $(4, 2)$ . Which of inner and dot runs on them in that order, and what shape comes out of the one that does?
Write np.outer(a, b) as an einsum string, then write $A^{\mathsf{T}}A$ as one. Check both against the plain versions.
A batch of RGB images has shape $(32, 3, 224, 224)$ . Which axes do you name to get the average brightness of each channel of each image, and what shape comes back?
Redraw Figure 3 from memory, labels included. The drawing is the test.