Tensors and shapes

What you will be able to do

Read the shape of every array of numbers that a network uses, before you run the code. A wrong shape either stops the code with an error, or, worse, gives a wrong answer with no error at all.

NotebookDo this lesson in a notebookColab · Kaggle · GitHub

The notebook has every code cell of this lesson, with the tasks and checks. It runs in Colab or Kaggle, where PyTorch is already installed, so you need no setup. About 30 minutes.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 50 minutes

What you need to know

• A matrix and its shape, one row per observation and one column per feature, from Represent observations in the mathematics course.
• The shape rule of matrix multiplication, (n × k) × (k × m) → (n × m), from Vector operations.
• A layer as many neurons that read the same inputs, from Neurons and layers earlier in this course.

What you need on your computer

• Nothing for the notebook: it opens in Colab or Kaggle with a free account. The next lesson, Set up PyTorch, installs PyTorch on your own computer, if you prefer to work there.

Tested with Python 3.12 and PyTorch 2.14.1 on the CPU (macOS).

Eggs, cartons and crates

Think of a shop that sells eggs. One egg is one thing. A row of 6 eggs is a line. A carton holds 2 rows of 6. A crate holds 10 cartons. To describe a full crate, you give one count for each direction: 10 cartons, 2 rows, 6 eggs. That is 10 × 2 × 6 = 120 eggs.

The number 120 alone is not enough. A delivery note that says “120 eggs” does not tell you how they are packed. They could be 20 cartons of 6, or 12 rows of 10. To unpack them correctly, you need the count in each direction, in the right order.

The problem: twelve numbers, two meanings

A neural network has the same problem with numbers. In the mathematics course, Priya described three tickets with four features each: the length, the refund words, the past tickets and a photo flag. That is 12 numbers. PyTorch, the library of this course, can hold them as one long list. reshape() cuts the list into rows. Here it cuts the same 12 numbers in two ways:

import torch

flat = torch.tensor([2.0, 2.0, 0.0, 1.0, 1.0, 0.0, 3.0, 0.0, 3.0, 1.0, 1.0, 1.0])   # 12 numbers
print(flat.reshape(3, 4))    # 3 tickets, 4 features each
print(flat.reshape(4, 3))    # 4 "tickets", 3 "features" each

The output:

tensor([[2., 2., 0., 1.],
        [1., 0., 3., 0.],
        [3., 1., 1., 1.]])
tensor([[2., 2., 0.],
        [1., 1., 0.],
        [3., 0., 3.],
        [1., 1., 1.]])

The first table is Priya’s: one row for each ticket. The second table has the same numbers, and no error appears. But it describes 4 tickets that do not exist. Its second row mixes the photo flag of ticket 1 with the first two features of ticket 2. A network that reads these rows learns from nonsense. So the network needs more than the numbers. It needs to know how they are arranged.

The idea: a tensor is a box of numbers with a shape

A tensor is a box of numbers, together with its shape: one count for each direction. Each direction is a dimension. The word comes from the Latin for “to stretch”. Physicists first used it for the forces that stretch a material. In deep learning, it means an array with any number of dimensions. PyTorch stores every input, weight and output of a network as a tensor.

Four tensors. The shape lists one count for each dimension, the outer dimension first.

Text description of the diagram

The diagram shows four boxes of numbers from left to right (top to bottom on a phone). A scalar has 0 dimensions: one cell with the score 5.5. Its shape is empty, written (). A vector has 1 dimension: one row of 4 cells, 2, 2, 0 and 1, the features of one ticket. Its shape is (4,). A matrix has 2 dimensions: 3 rows of 4 cells, the three tickets. Its shape is (3, 4). An arrow down the rows is labelled dim 0, and an arrow across the columns is labelled dim 1. A 3-D tensor has 3 dimensions: two 3 by 4 grids, one behind the other. Its shape is (2, 3, 4): 2 tickets, 3 words each, 4 numbers for each word.

Here are the same four tensors in PyTorch. torch.tensor() makes a tensor from Python lists, and torch.zeros() makes one of the given shape, filled with zeros. t.ndim is the number of dimensions, t.shape is the shape, and t.numel() counts the numbers.

shapes.py

import torch

score = torch.tensor(5.5)                      # one number
ticket = torch.tensor([2.0, 2.0, 0.0, 1.0])    # one ticket: 4 features
batch = torch.tensor([[2.0, 2.0, 0.0, 1.0],
                      [1.0, 0.0, 3.0, 0.0],
                      [3.0, 1.0, 1.0, 1.0]])   # 3 tickets, 4 features each
words = torch.zeros(2, 3, 4)                   # 2 tickets, 3 words, 4 numbers per word

for name, t in [("score", score), ("ticket", ticket), ("batch", batch), ("words", words)]:
    print(f"{name:<7} ndim {t.ndim}   shape {tuple(t.shape)}   numbers {t.numel()}")

The output:

score   ndim 0   shape ()   numbers 1
ticket  ndim 1   shape (4,)   numbers 4
batch   ndim 2   shape (3, 4)   numbers 12
words   ndim 3   shape (2, 3, 4)   numbers 24

The count of numbers is always the sizes multiplied together: 3 × 4 = 12 and 2 × 3 × 4 = 24. The comma in (4,) is how Python writes a tuple with one value. PyTorch prints a shape as torch.Size, which works like a tuple of whole numbers:

print(batch.shape)
print(batch.shape[0], "tickets,", batch.shape[1], "features")

The output:

torch.Size([3, 4])
3 tickets, 4 features

The names scalar, vector and matrix are the special cases with 0, 1 and 2 dimensions. You met a 3-D tensor in the mathematics course as a colour photo: height × width × 3 colours. In Module 5 of this course, each word of a ticket becomes a row of numbers. Then a batch of tickets is a 3-D tensor like words.

The batch dimension

In batch, each row is one ticket and each column is one feature. The first dimension, dim 0, counts the tickets. A group of examples that a network handles together is a batch, so dim 0 is the batch dimension. The last dimension, dim 1, holds the features of each example. This is the layout of almost every network in this course: (batch, features).

Why handle tickets in a batch, and not one by one? One matrix product scores all the tickets at once. A computer does that much faster than many small products, and the code stays the same for 3 tickets or 3,000.

Take care when you pick one ticket out of a batch. An index removes the batch dimension, and a slice keeps it:

print(batch[0], batch[0].shape)
print(batch[0:1], batch[0:1].shape)

The output:

tensor([2., 2., 0., 1.]) torch.Size([4])
tensor([[2., 2., 0., 1.]]) torch.Size([1, 4])

Both hold the same four numbers. But batch[0:1] is still a batch: a batch of one ticket. Code that expects (batch, features) works with it, as The forward pass shows.

Multiply: the shape rule you know

A layer of a network computes weighted sums, as in Neurons and layers. In PyTorch, @ is matrix multiplication. Here are Priya’s urgency weights from the mathematics course, one weight for each feature:

w = torch.tensor([0.5, 2.0, 1.0, 0.5])   # one weight per feature
scores = batch @ w
print(scores, scores.shape)

The output:

tensor([5.5000, 3.5000, 5.0000]) torch.Size([3])

The first ticket’s score is 2 × 0.5 + 2 × 2 + 0 × 1 + 1 × 0.5 = 5.5. The rule from the mathematics course holds: the inner numbers must be equal, and the outer numbers give the result. Now Tomás wants one score for each of the 5 teams. He writes the weights with one row for each team:

W = torch.zeros(4, 5)          # 4 features in, 5 teams out
print((batch @ W).shape)
W_wrong = torch.zeros(5, 4)    # one row per team
print((batch @ W_wrong).shape)

The output:

torch.Size([3, 5])
RuntimeError: mat1 and mat2 shapes cannot be multiplied (3x4 and 5x4)

PyTorch calls the two matrices mat1 and mat2, and gives both shapes. The inner numbers are 4 and 5, so there is no product. W_wrong.T turns the matrix, so its shape becomes (4, 5), and the product works. Read every shape error in this way. Find the two shapes, then find the pair of numbers that must be equal.

Broadcasting: one row for every ticket

A shop writes “delivery: €5 per order” once, at the top of its price list. Nobody copies €5 into every row of the order table. Everybody understands that the same fee applies to every order.

Priya wants the same thing with tensors. She wants to know how far each ticket is from the average ticket. So she computes the mean of each feature, a row of 4 numbers, and subtracts it from every ticket. Without help, she must first copy the row 3 times, to get a (3, 4) tensor of means. With 4,696 tickets, that copy has 4,696 rows.

PyTorch does the copying for her, without storing the copies. This is broadcasting: the smaller tensor is “broadcast” to the larger one, like a radio signal that reaches every house. The rule compares the two shapes from the right, one dimension at a time. Two sizes fit when they are equal, or when one of them is 1, or when one of them is missing. A size of 1, or a missing dimension, is stretched to match.

The broadcasting rule on three cases. A missing dimension or a 1 stretches; any other pair must be equal.

Text description of the diagram

The diagram shows three cases side by side (one under the other on a phone). Each case puts the shape of batch, 3 and 4, in two boxes, above the shape of a second tensor, aligned on the right. In the first case, batch minus mean, mean has an empty, dashed box on the left and 4 on the right. 4 equals 4, and the missing dimension stretches to 3. The result is (3, 4). In the second case, batch divided by totals, totals has 3 and 1. The 1 stretches to 4, and 3 equals 3. The result is (3, 4). In the third case, batch divided by totals_flat, totals_flat has an empty box and 3. Compared from the right, 4 and 3 are not equal, and neither is 1. The result is an error.

Here is the first case. batch.mean(dim=0) takes the mean down dim 0, over the tickets, so one number is left for each feature:

mean = batch.mean(dim=0)     # the mean of each column: shape (4,)
print(mean)
print(batch - mean)          # (3, 4) - (4,)

The output:

tensor([2.0000, 1.0000, 1.3333, 0.6667])
tensor([[ 0.0000,  1.0000, -1.3333,  0.3333],
        [-1.0000, -1.0000,  1.6667, -0.6667],
        [ 1.0000,  0.0000, -0.3333,  0.3333]])

Check one number. Ticket 2 has 3 past tickets, and the mean is 1.3333, so 3 − 1.3333 = 1.6667. Broadcasting also works across the columns. Here each ticket’s numbers are divided by that ticket’s total. keepdim=True keeps the summed dimension with a size of 1, so the totals form a column:

totals = batch.sum(dim=1, keepdim=True)   # one total per ticket: shape (3, 1)
print(totals)
print(batch / totals)                     # (3, 4) / (3, 1)

The output:

tensor([[5.],
        [4.],
        [6.]])
tensor([[0.4000, 0.4000, 0.0000, 0.2000],
        [0.2500, 0.0000, 0.7500, 0.0000],
        [0.5000, 0.1667, 0.1667, 0.1667]])

Without keepdim=True, the totals have the shape (3,). Compared from the right, 4 and 3 do not fit:

totals_flat = batch.sum(dim=1)            # shape (3,): keepdim forgotten
print(totals_flat.shape)
print(batch / totals_flat)

The output:

torch.Size([3])
RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 1

A “non-singleton” dimension is one whose size is not 1. The message names the sizes, 4 and 3, and the dimension where they meet, dim 1. Predict first, then try each operation in the explorer. In the explorer, W is a smaller matrix of shape (4, 2), two teams instead of five, so that all its numbers fit on the screen.

Explore

Shapes in, shape out

Shape of batch

(3, 4)

Shape of mean

(4,)

Result

(3, 4)

Result: shape (3, 4). Compare the shapes from the right: 4 and 4 are equal. mean has no first dimension, so PyTorch uses its one row for all 3 tickets (the light copies). Each ticket now shows how far it is from the average ticket.

batch

(3, 4) float32

2	2	0	1
1	0	3	0
3	1	1	1

−

mean

(4,) float32

stretched to (3, 4)

2	1	1.33	0.67
2	1	1.33	0.67
2	1	1.33	0.67

→

result

(3, 4) float32

0	1	−1.33	0.33
−1	−1	1.67	−0.67
1	0	−0.33	0.33

PyTorch prints

>>> (batch - mean).shape
torch.Size([3, 4])

Real results and error messages from PyTorch 2.14.1 on the CPU. Light, dashed cells are the copies that broadcasting uses but does not store.

Text description of the explorer

The explorer has one row of buttons, one for each operation. The operations are batch − mean, batch / totals, batch / totals_flat, batch @ w, batch @ W, batch @ W.T, scores − targets and counts @ w. For each one, it shows the shapes that go in, the tensors with their values, the result, and what PyTorch prints.

batch − mean: (3, 4) and (4,) give (3, 4); the row of means is used for all 3 tickets. batch / totals: (3, 4) and (3, 1) give (3, 4); each row adds up to 1. batch / totals_flat: (3, 4) and (3,) stop with an error. The message is “The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 1”. batch @ w: (3, 4) and (4,) give (3,), the scores 5.5, 3.5 and 5.0. batch @ W, with W of shape (4, 2): the result is (3, 2). batch @ W.T, with W.T of shape (2, 4), stops with “mat1 and mat2 shapes cannot be multiplied (3x4 and 2x4)”. scores − targets: (3,) and (3, 1) give (3, 3) with no error, a wrong answer that the common mistake below explains. counts @ w: the shapes fit, but counts holds whole numbers and w holds decimals. PyTorch stops with “dot : expected both vectors to have same dtype, but found Long and Float”.

When should you rely on broadcasting? Use it for the operations that you mean: one row of values for every example, or one column of values for every feature. Its cost is that it never asks if you meant it. The common mistake at the end of this lesson shows a shape that broadcasts by accident.

The type of the numbers: dtype

Each tensor also has a dtype, the type of its numbers. torch.tensor() makes whole numbers int64 and decimals float32, which uses 4 bytes for each number. Networks compute with float32. A product does not mix the two types:

counts = torch.tensor([2, 2, 0, 1])     # whole numbers
print(counts.dtype, w.dtype)
print(counts @ w)

The output:

torch.int64 torch.float32
RuntimeError: dot : expected both vectors to have same dtype, but found Long and Float

PyTorch calls int64 “Long” and float32 “Float” in its messages. The fix is counts.float() @ w, which gives tensor(5.5000). Word counts are whole numbers, so The forward pass makes them as decimals from the start.

Where the numbers live: the device

A tensor also has a device: the hardware whose memory holds it. In this course, every tensor lives on the CPU, and print(batch.device) shows cpu. A tensor can move to a graphics card with .to("cuda"), but all tensors in one operation must be on the same device. The next lesson shows how to check for a graphics card. You do not need one.

NumPy and PyTorch

In the machine learning course, your arrays were NumPy arrays. torch.from_numpy() turns a NumPy array into a tensor, and .numpy() turns a tensor back. Both share the same memory, so nothing is copied. Watch the dtype: NumPy makes decimals float64, 8 bytes each, not float32.

import numpy as np

from_numpy = torch.from_numpy(np.array([2.0, 2.0, 0.0, 1.0]))
print(from_numpy.dtype)
print(from_numpy @ w)

The output:

torch.float64
RuntimeError: dot : expected both vectors to have same dtype, but found Double and Float

“Double” is float64. Write from_numpy.float() @ w, and the result is tensor(5.5000). In the other direction, (batch @ w).numpy() gives the NumPy array [5.5 3.5 5. ].

The Larkfield tickets as tensors

Now the real task. In The forward pass, each ticket becomes 2,864 numbers, one for each word of a fixed list. A batch of 8 tickets is then a tensor of shape (8, 2864). The first layer of the network turns 2,864 numbers into 64. Before you run any code, read the shapes: (8, 2864) @ (2864, 64). The inner numbers match, so the result is (8, 64).

tickets = torch.zeros(8, 2864)       # 8 tickets, 2,864 word counts each
W1 = torch.zeros(2864, 64)           # 2,864 inputs, 64 outputs
print((tickets @ W1).shape)
print(tickets.element_size(), "bytes per number,", tickets.element_size() * tickets.numel(), "bytes in all")
print("all 4,696 training tickets:", 4696 * 2864 * 4, "bytes")

The output:

torch.Size([8, 64])
4 bytes per number, 91648 bytes in all
all 4,696 training tickets: 53797376 bytes

The shape also tells you the memory. Each float32 number takes 4 bytes, so the batch of 8 takes 91,648 bytes. All 4,696 training tickets take about 54 MB. That fits easily on a laptop. But the memory grows with every dimension. Amira will run the network in the help-desk app, so she asks about memory for every model in this course.

When to read shapes, and what it costs

Read the shapes before you run code, every time two tensors meet: in a product, in +, −, * or /, and when a tensor goes into a layer. Write the expected shape in a comment, as the code of this lesson does, and print x.shape when you are not sure. This habit costs a few seconds. It saves the hours that a wrong shape costs when it raises no error. Shapes do not tell you everything. Two tensors with the right shapes can still hold the wrong numbers, as the twelve numbers at the start showed.

Guided practice: shapes in the notebook

1. Open the notebook from the box Do this lesson in a notebook at the top of this page.
2. Run the setup cells. The first one prints a line that starts with PyTorch 2..
3. In section 1, make the four tensors. You see the shapes (), (4,), (3, 4) and (2, 3, 4).
4. In section 2, predict the shape of batch[0] and of batch[0:1]. Then run the cell. You see torch.Size([4]) and torch.Size([1, 4]).
5. In section 3, compute scores = batch @ w. You see tensor([5.5000, 3.5000, 5.0000]).
6. In section 4, predict which of the three broadcasting cells raises an error. Then run them. Only batch / totals_flat raises one.
7. In the Your turn cell, compute centered: each ticket minus the mean of each feature. Then run the check cell. It prints Check passed.

Show an example answer to step 7

centered = batch - batch.mean(dim=0)
print(centered.shape)

torch.Size([3, 4])

batch.mean(dim=0) has the shape (4,), one mean for each feature. Broadcasting uses it for each of the 3 tickets, so the result has the shape of batch.

Your turn: five tickets

Change one thing: the batch now has 5 tickets, five = torch.ones(5, 4). Before you run anything, write down the shape of five @ w, of five - five.mean(dim=0) and of five / five.sum(dim=1, keepdim=True). Then check your answers in the notebook.

Show the answer

five = torch.ones(5, 4)                       # 5 tickets, 4 features
print((five @ w).shape)
print((five - five.mean(dim=0)).shape)
print((five / five.sum(dim=1, keepdim=True)).shape)

torch.Size([5])
torch.Size([5, 4])
torch.Size([5, 4])

Only the batch dimension changed, from 3 to 5. The weights, the means and the features still have 4 numbers, so the same code works. This is why the batch comes first: the code does not depend on how many tickets there are.

Common mistake

The error that does not appear

Symptom: Priya compares the three scores with the scores that Grace expected, 5, 4 and 5. She computes the mean absolute error, from Regression metrics in the machine learning course. She expects about 0.3. She gets this, and no error:

targets = torch.tensor([[5.0], [4.0], [5.0]])   # one column: shape (3, 1)
errors = scores - targets                       # (3,) - (3, 1)
print(errors.shape)
print(errors)
print("mean absolute error:", round(errors.abs().mean().item(), 3))

torch.Size([3, 3])
tensor([[ 0.5000, -1.5000,  0.0000],
        [ 1.5000, -0.5000,  1.0000],
        [ 0.5000, -1.5000,  0.0000]])
mean absolute error: 0.778

Cause: scores has the shape (3,) and targets has the shape (3, 1). Compared from the right, 3 meets 1, and then nothing meets 3. Both tensors stretch, so PyTorch subtracts every target from every score: 9 differences instead of 3. The mean of those 9 is not the error of the tickets.

Fix: make the two shapes equal before you subtract. targets.squeeze(1) removes the dimension of size 1:

targets = targets.squeeze(1)    # remove the dimension of size 1: shape (3,)
errors = scores - targets
print(errors.shape, errors)
print("mean absolute error:", round(errors.abs().mean().item(), 3))

torch.Size([3]) tensor([ 0.5000, -0.5000,  0.0000])
mean absolute error: 0.333

What it means for a real AI system: a shape that broadcasts by accident gives a number that looks normal. 0.778 is a believable error. Nothing crashes, and the wrong number goes into a report or into training. Print the shape of every result that you did not expect to change, and check that it is the shape in your comment.

See it end to endClassify clothing imagesCase study

The clothing case study turns real photos into a tensor of shape (N, 1, 28, 28), and shows the error when the channel dimension is missing.

Knowledge check

Check your understanding

A tensor has the shape (16, 2864). Say in words what each number means for a batch of tickets, how many numbers it holds, and how many bytes they take as float32.

Show the answer

16 is the batch dimension: 16 tickets. 2,864 is the feature dimension: 2,864 numbers for each ticket, one for each entry of the word list. The tensor holds 16 × 2,864 = 45,824 numbers. Each float32 number takes 4 bytes, so they take 183,296 bytes, about 0.18 MB.

Check your understanding

Predict the result shape, or an error, for each pair: (8, 2864) @ (2864, 64); (8, 64) + (64,); (8, 64) + (8,); (8, 5) - (8, 1).

Show the answer

• (8, 2864) @ (2864, 64) → (8, 64): the inner numbers are equal.
• (8, 64) + (64,) → (8, 64): from the right, 64 = 64, and the missing dimension stretches to 8. This is how a layer adds its bias to every ticket.
• (8, 64) + (8,) → error: from the right, 64 meets 8. They are not equal, and neither is 1.
• (8, 5) − (8, 1) → (8, 5): from the right, 1 stretches to 5, and 8 = 8. One number per ticket is subtracted from each of its 5 scores.

Check your understanding

Tomás says: 'If the shapes were wrong, PyTorch would tell me.' Use this lesson to explain when that is true and when it is not.

Show the answer

It is true when no rule can make the shapes fit. Examples are a product with different inner numbers, or two sizes that differ when neither is 1. Then PyTorch stops and names both shapes. It is not true when broadcasting can stretch a tensor. Scores of shape (3,) minus targets of shape (3, 1) give a (3, 3) table with no error. The mean absolute error is then 0.778 instead of 0.333. PyTorch checks that the shapes fit together. It cannot check that they mean what you intended. So read the shapes before you run the code, and print them after.

Recap

• A tensor is a box of numbers with a shape: one count for each dimension. A scalar has shape (), a vector (4,), a matrix (3, 4), a 3-D tensor (2, 3, 4).
• In a network, the first dimension is usually the batch: (batch, features). batch[0] drops it; batch[0:1] keeps it.
• @ multiplies matrices: the inner numbers must be equal. PyTorch’s message names both shapes: “mat1 and mat2 shapes cannot be multiplied (3x4 and 5x4)”.
• Broadcasting compares shapes from the right. Equal sizes, a 1 or a missing dimension fit; a 1 or a missing dimension stretches. It can also stretch by accident, with no error: (3,) − (3, 1) gives (3, 3).
• Each tensor has a dtype (float32 for networks, int64 for whole numbers, float64 from NumPy) and a device (the CPU in this course). An operation needs the same dtype and device for all its tensors.

Next, you set up PyTorch on your own computer, or check it in Colab. Then you run a first program and compare its output, number by number.

Next: Set up PyTorch