PT01: Introduction to PyTorch (instructor version)

This notebook covers the PyTorch we need for the first labs: tensors, batched linear models, gradients and parameter updates. If you have used PyTorch before, skim Sections 1 to 5 and try the short check at the end of Section 5. If you have not, read and run everything up to that check before PT02 (the PyTorch documentation is a good companion). Knowing NumPy helps, since most of the array syntax is the same.

Everything after the check is optional: a few short tensor exercises, a closer look at autograd with some debugging cases, and how tensors are stored in memory. You can come back to these later.

A good habit: before running a cell, predict its output, or at least its shape. If something surprises you, change the example and run it again.

This is the instructor notebook. The worked examples are shared with the student version; answers to the check and the exercises appear only here. Edit this file and rebuild the student notebook with make_student_notebook.py.

We annotate the shapes of function inputs and outputs with jaxtyping. For example, Float[torch.Tensor, "batch features"] is a floating-point tensor with two axes, named batch and features. The annotations are only documentation: nothing checks them when the code runs. The next cell installs the package, which you need on Colab.

%pip install jaxtyping

Requirement already satisfied: jaxtyping in /Users/simonescardapane/Repositories/teaching-materials/build/nb-venv/lib/python3.12/site-packages (0.3.11)
Requirement already satisfied: wadler-lindig>=0.1.3 in /Users/simonescardapane/Repositories/teaching-materials/build/nb-venv/lib/python3.12/site-packages (from jaxtyping) (0.1.7)

[notice] A new release of pip is available: 24.0 -> 26.2.1
[notice] To update, run: /Users/simonescardapane/Repositories/teaching-materials/build/nb-venv/bin/python -m pip install --upgrade pip

Note: you may need to restart the kernel to use updated packages.

import torch
from jaxtyping import Float, Int
from typing import Optional
import torch.nn.functional as F

torch.manual_seed(7)
torch.set_printoptions(precision=3, sci_mode=False)
print("PyTorch", torch.__version__)

PyTorch 2.14.1

1. Tensors

The basic data type is torch.Tensor, PyTorch's version of a NumPy array. We can build one from a list with torch.tensor, or with functions such as torch.zeros and torch.randn. In this course a batch of data has one example per row and one feature per column. Here $X \sim (B,D)$, with $B=3$ examples and $D=2$ features.

X = torch.tensor([[1., 2.], [3., 1.], [2., 4.]])
print(X)
print("Shape:", X.shape)
print("Dtype:", X.dtype)
print("Device:", X.device)

tensor([[1., 2.],
        [3., 1.],
        [2., 4.]])
Shape: torch.Size([3, 2])
Dtype: torch.float32
Device: cpu

Every tensor also lives on a device (CPU or GPU). X.to(...) moves it to another one; we do not need this yet.

Indexing works as in NumPy. Look at the last two lines: indexing a single column drops that axis, while slicing keeps it.

print("Second example:", X[1])
print("First two examples:\n", X[:2])
print("Last feature:", X[:, -1], X[:, -1].shape)
print("Last feature, keeping its axis:\n", X[:, -1:], X[:, -1:].shape)

Second example: tensor([3., 1.])
First two examples:
 tensor([[1., 2.],
        [3., 1.]])
Last feature: tensor([2., 1., 4.]) torch.Size([3])
Last feature, keeping its axis:
 tensor([[2.],
        [1.],
        [4.]]) torch.Size([3, 1])

2. A batch of linear predictions

For a single example $x$, a linear model with $C$ outputs computes $z=Wx+b$. With the examples stored as rows, we can process the whole batch at once:

$$Z=XW^\top+b.$$

Here $W\sim(C,D)$, $b\sim(C,)$ and $Z\sim(B,C)$. The bias is a single vector, and PyTorch broadcasts it: it is added to every row. @ is matrix multiplication, while * is elementwise multiplication. The last line checks the first row of Z against the single-example formula.

W = torch.tensor([[1., -1.], [0.5, 2.]])  # (C, D)
b = torch.tensor([0.2, -0.3])             # (C,)
Z = X @ W.T + b                         # (B, C)
print(Z)
torch.testing.assert_close(Z[0], W @ X[0] + b)

tensor([[-0.800,  4.200],
        [ 2.200,  3.200],
        [-1.800,  8.700]])

Reductions such as sum and mean remove the axis they reduce over. With keepdim=True the axis stays, with size one, which is what we need when the result is broadcast back against the original tensor.

For example, to center each feature we average over the examples (dim=0), not over the features. Change dim=0 to dim=1 below: what gets centered now?

feature_mean = X.mean(dim=0, keepdim=True)
centered = X - feature_mean
print("Mean shape:", feature_mean.shape)
print(centered)
torch.testing.assert_close(centered.mean(dim=0), torch.zeros(2), atol=1e-6, rtol=0)

Mean shape: torch.Size([1, 2])
tensor([[-1.000, -0.333],
        [ 1.000, -1.333],
        [ 0.000,  1.667]])

3. Computing gradients

Take a regression model $\hat y=Xw+b$ with a mean squared error loss. We mark the tensors we want gradients for (the parameters) with requires_grad=True. Before running the next cell, predict the shapes of the loss, w.grad and bias.grad.

y = torch.tensor([1., -1., 2.])
w = torch.tensor([0.2, -0.1], requires_grad=True)
bias = torch.tensor(0., requires_grad=True)

prediction = X @ w + bias
loss = ((prediction - y) ** 2).mean()
print("Loss:", loss.item(), "Shape:", loss.shape)
print("Parameter is a leaf:", w.is_leaf)
print("Prediction is a leaf:", prediction.is_leaf)
loss.backward()
print("w.grad:", w.grad)
print("bias.grad:", bias.grad)

Loss: 2.4166667461395264 Shape: torch.Size([])
Parameter is a leaf: True
Prediction is a leaf: False
w.grad: tensor([-0.333, -5.667])
bias.grad: tensor(-1.)

Calling backward() on the scalar loss computes its gradient with respect to the tensors we created with requires_grad=True, and stores it in their .grad field. These tensors are called leaves. Intermediate results such as prediction take part in the computation too, but their gradients are not stored.

During the forward pass, PyTorch records the operations it runs; backward() goes back through this record. Afterwards most of the record is freed, so each new gradient needs a new forward pass. We will see how this works in detail when we study automatic differentiation.

4. Updating the parameters

A gradient descent step is $\theta\leftarrow\theta-\eta\nabla L(\theta)$. We run it inside torch.no_grad(): the update is not part of the model, so autograd should not record it (and PyTorch refuses to modify a leaf that requires gradients in place outside such a block). Then we clear the gradients, because backward() adds to .grad instead of overwriting it.

lr = 0.01
before = loss.item()
with torch.no_grad():
    w -= lr * w.grad
    bias -= lr * bias.grad
w.grad = None
bias.grad = None

after = ((X @ w + bias - y) ** 2).mean()
print(f"Loss before: {before:.4f}; after one update: {after.item():.4f}")
assert after.item() < before  # This example uses a sufficiently small step.

Loss before: 2.4167; after one update: 2.1115

We can see this accumulation without training anything. Keep a fixed and differentiate the same function twice, with a new forward pass each time:

a = torch.tensor([1., 2.], requires_grad=True)
a.square().sum().backward()
first_gradient = a.grad.clone()
a.square().sum().backward()
print("First backward:", first_gradient)
print("Second backward:", a.grad)
torch.testing.assert_close(a.grad, 2 * first_gradient)
a.grad = None

First backward: tensor([2., 4.])
Second backward: tensor([4., 8.])

To log values during training, use loss.item(), which returns a plain Python number. If you need to keep a tensor, tensor.detach() returns one that is cut off from the recorded computation.

torch.autograd.grad is a second way to call autograd: it returns the gradients instead of adding them to .grad. We will use it again when we study automatic differentiation.

a = torch.tensor([1., 2.], requires_grad=True)
(gradient,) = torch.autograd.grad(a.square().sum(), a)
print(gradient)
assert a.grad is None

tensor([2., 4.])

5. Classification uses the same steps

For classification with $C$ classes, the model outputs $Z\sim(B,C)$. These unnormalized scores are called logits. The targets $y\sim(B,)$ are integer class indices. A softmax over each row turns the logits into probabilities, which is useful when we want to look at them.

F.cross_entropy takes the logits, not the probabilities. It applies the softmax and the logarithm internally, in a numerically stable way.

logits = torch.tensor([[2., 0., -1.], [-1., 0., 2.]], requires_grad=True)
labels = torch.tensor([0, 2], dtype=torch.long)
classification_loss = F.cross_entropy(logits, labels)
classification_loss.backward()
print("Probabilities:\n", logits.detach().softmax(dim=-1))
print("Loss:", classification_loss.item())
print("Gradient shape:", logits.grad.shape)

Probabilities:
 tensor([[0.844, 0.114, 0.042],
        [0.042, 0.114, 0.844]])
Loss: 0.1698460429906845
Gradient shape: torch.Size([2, 3])

Every training step has the same structure: clear the gradients, compute the predictions and a scalar loss, call backward(), update the parameters. For evaluation we only need the predictions, so there is no reason to record anything.

Two tools are easy to confuse here. model.eval() changes how some layers behave, such as dropout and batch normalization, which we will see later. torch.no_grad() turns off recording. In PT02 we use both when evaluating a model.

This is the end of the required part. Before PT02, make sure you can explain how the bias is broadcast, what shape each gradient has, and what happens if we forget to clear the gradients.

Before PT02: a short check

Answer these without running code first, then use the next cell to test your answers. The examples above contain the ideas, but not the answers.

1. A has shape (3, 2) and we want to center each row. A - A.mean(dim=1) raises an error: why, and what is the fix? Now take S = torch.tensor([[1., 3.], [5., 7.]]). Does S - S.mean(dim=1) raise an error? If not, what does it compute, and how does it differ from the correctly centered rows?
2. a = torch.tensor([1., 2.], requires_grad=True). We call a.square().sum().backward(), then (3 * a).sum().backward(), without clearing. What is a.grad, and of which function is it the gradient? We then update a -= 0.1 * a.grad inside no_grad and call a.square().sum().backward() once more, still without clearing. What are a and a.grad now, and what should a.grad have been?
3. For X @ W.T + b with shapes (5, 4), (3, 4) and (3,), what are the shapes of the logits, of the mean cross-entropy loss, and of the two parameter gradients? What happens if we call backward() on F.cross_entropy(logits, y, reduction="none") instead, and how can we get the same gradients as the mean loss from it?

If one of these is unclear, go back to the corresponding example before PT02.

# Try your repairs and shape predictions here.

Check: answers

1. Reducing over dim=1 gives shape (3,), which does not broadcast against (3, 2). Use A.mean(dim=1, keepdim=True), with shape (3, 1). With the square S, the row means [2, 6] have shape (2,) and broadcast as a row: the first mean is subtracted from the first column and the second from the second column. The result [[-1, -3], [3, 1]] has the right shape and no error, but the correct answer is [[-1, 1], [-1, 1]]. Square shapes hide axis mistakes.
2. [2, 4] + [3, 3] = [5, 7], the gradient of a.square().sum() + (3 * a).sum(): backward() adds, so the stored gradient belongs to the sum of the two functions. After the update a = [0.5, 1.3], and the third call adds [1.0, 2.6], giving [6.0, 9.6] instead of [1.0, 2.6]. Set a.grad = None before each backward call whose gradient we want on its own.
3. Logits (5, 3); mean loss (); weight gradient (3, 4); bias gradient (3,). With reduction="none" the loss has shape (5,) and backward() raises an error, because it needs a scalar to start from. Calling .mean() first gives the same gradients as before.

# Executable answers to the check.
S = torch.tensor([[1., 3.], [5., 7.]])
torch.testing.assert_close(S - S.mean(dim=1), torch.tensor([[-1., -3.], [3., 1.]]))
torch.testing.assert_close(S - S.mean(dim=1, keepdim=True), torch.tensor([[-1., 1.], [-1., 1.]]))

a = torch.tensor([1., 2.], requires_grad=True)
a.square().sum().backward()
(3 * a).sum().backward()
torch.testing.assert_close(a.grad, torch.tensor([5., 7.]))
with torch.no_grad():
    a -= 0.1 * a.grad
a.square().sum().backward()
torch.testing.assert_close(a, torch.tensor([0.5, 1.3]))
torch.testing.assert_close(a.grad, torch.tensor([6.0, 9.6]))

check_X, check_y = torch.randn(5, 4), torch.tensor([0, 2, 1, 1, 0])
check_W = torch.randn(3, 4, requires_grad=True)
check_b = torch.zeros(3, requires_grad=True)
per_example = F.cross_entropy(check_X @ check_W.T + check_b, check_y, reduction="none")
assert per_example.shape == (5,)
try:
    per_example.backward()
except RuntimeError as error:
    print("Non-scalar backward:", str(error).splitlines()[0])
else:
    raise AssertionError("backward() on a vector loss should fail")
per_example.mean().backward()
assert check_W.grad.shape == (3, 4) and check_b.grad.shape == (3,)
mean_W_grad, mean_b_grad = torch.autograd.grad(
    F.cross_entropy(check_X @ check_W.T + check_b, check_y), (check_W, check_b))
torch.testing.assert_close(check_W.grad, mean_W_grad)
torch.testing.assert_close(check_b.grad, mean_b_grad)

Non-scalar backward: grad can be implicitly created only for scalar outputs

Optional: tensor semantics challenges

Three short exercises that only need Sections 1 and 2, plus the softmax from Section 5. Each contains an expression that runs without errors but does not compute what we want. You can do them now or come back to them during a lab.

For each one, predict the shape and values of the result before running the cell. Then explain the mistake and fix it. Finally, write assertions on the shape and on values you computed by hand, and check that at least one of them fails on the original expression.

1. One error per example

We have three predictions stored as a column, with shape (3, 1), and three targets with shape (3,). We want the three prediction errors and their mean squared value.

What does each entry of suspect_errors compare? Fix the subtraction so that there is one error per example, in a tensor of shape (3,). Would checking that the loss is a scalar catch this bug?

predictions = torch.tensor([[1.], [2.], [4.]])  # (B, 1)
targets = torch.tensor([0., 2., 5.])            # (B,)
suspect_errors = predictions - targets
suspect_mse = suspect_errors.square().mean()
print(suspect_errors.shape, suspect_errors)
print("MSE:", suspect_mse)

# Write the corrected errors and MSE, then add your assertions.

torch.Size([3, 3]) tensor([[ 1., -1., -4.],
        [ 2.,  0., -3.],
        [ 4.,  2., -1.]])
MSE: tensor(5.778)

2. Which axis contains the examples?

Each row of batch is an example with two features. The model predicts the first feature plus ten times the second, giving one prediction per row.

The batch is square, so both batch @ weights and batch.T @ weights run. Predict suspect_predictions, then fix it and check your answer. Why can a shape assertion not catch this mistake? As a second test, add a third example [5., 6.] and predict what happens to each expression.

batch = torch.tensor([[1., 2.], [3., 4.]])  # (B, D)
weights = torch.tensor([1., 10.])           # (D,)
suspect_predictions = batch.T @ weights
print(suspect_predictions.shape, suspect_predictions)

# Compute one prediction per row, then add your assertions.
# Repeat with a third example.

torch.Size([2]) tensor([31., 42.])

3. One distribution per example

scores contains logits for three examples (rows) and three classes (columns). We want, for each example, a probability distribution over the classes.

Predict which sums of suspect_probabilities equal one: the rows or the columns? Fix it and check that each example's probabilities sum to one. Why can neither a shape check nor a check that "all entries are between 0 and 1" catch this? Does the bug change any predicted class?

scores = torch.tensor([[3., 1., 0.], [0., 0., 1.], [4., 0., 0.]])  # (B, C)
suspect_probabilities = scores.softmax(dim=0)
print(suspect_probabilities.shape, suspect_probabilities)
print("Predicted classes:", suspect_probabilities.argmax(dim=1))

# Normalize over the classes, then add your assertions.

torch.Size([3, 3]) tensor([[0.265, 0.576, 0.212],
        [0.013, 0.212, 0.576],
        [0.721, 0.212, 0.212]])
Predicted classes: tensor([1, 2, 0])

Tensor semantics: solutions

1. One error per example. Broadcasting produces a (3, 3) matrix: entry (i, j) compares prediction i with target j. Its rows are [1, -1, -4], [2, 0, -3], and [4, 2, -1]. Averaging all nine squared entries gives 52/9, but the intended errors are [1, 0, -1], with MSE 2/3. Both losses are scalars.

errors = predictions.squeeze(-1) - targets
mse = errors.square().mean()
assert errors.shape == targets.shape == (3,)
torch.testing.assert_close(errors, torch.tensor([1., 0., -1.]))
torch.testing.assert_close(mse, torch.tensor(2. / 3.))

The shape assertion rejects suspect_errors; the loss-value assertion rejects suspect_mse. squeeze(-1) removes only the last axis, so a batch of size one still has its batch axis.

2. Which axis contains the examples? The original computes [31, 42] by combining different examples. Its shape (2,) looks right. The intended predictions are [21, 43].

row_predictions = batch @ weights
assert row_predictions.shape == (batch.shape[0],)
torch.testing.assert_close(row_predictions, torch.tensor([21., 43.]))

larger_batch = torch.cat([batch, torch.tensor([[5., 6.]])], dim=0)
larger_predictions = larger_batch @ weights
assert larger_predictions.shape == (3,)
torch.testing.assert_close(larger_predictions, torch.tensor([21., 43., 65.]))

The value assertion rejects suspect_predictions. With three examples, larger_batch.T @ weights raises a dimension mismatch: the last axis of the transposed batch has size three, while weights has size two.

3. One distribution per example. softmax(dim=0) normalizes each column, so the three examples compete for each class. The columns sum to one, while the rows sum to about 1.05, 0.80 and 1.15. Every entry is still between 0 and 1. The first example is predicted as class 1 (probability 0.58) although its largest logit is class 0: normalizing over the classes gives [0.84, 0.11, 0.04].

probabilities = scores.softmax(dim=-1)
assert probabilities.shape == scores.shape
torch.testing.assert_close(probabilities.sum(dim=-1), torch.ones(3))
torch.testing.assert_close(probabilities.argmax(dim=-1), scores.argmax(dim=-1))
torch.testing.assert_close(probabilities.argmax(dim=-1), torch.tensor([0, 2, 0]))

The row-sum assertion rejects suspect_probabilities, and so does the comparison of predicted classes. dim=-1 names the class axis even if we later add more leading axes. For predictions alone the softmax is unnecessary: scores.argmax(dim=-1) gives the same classes, which is why PT02 keeps logits and applies a softmax only when it needs probabilities.

Optional: a closer look at autograd

This section goes deeper into Sections 3 and 4: what autograd stores on each tensor, and how to cut a tensor off from the recorded computation. The autograd detective cases after it build on this material.

What a tensor knows about autograd

A few attributes tell us how autograd sees a tensor:

Attribute	Meaning
requires_grad	Autograd records the operations that use this tensor.
is_leaf	No recorded operation produced this tensor. Parameters we create ourselves are leaves. Tensors that do not require gradients are also leaves.
grad_fn	For a tensor produced by a recorded operation, the backward step of that operation. None for leaves. It is not a gradient.
grad	Where backward() stores the gradient. It starts as None and, by default, is filled only for leaves.
retains_grad	Whether an intermediate tensor was asked to keep its gradient with retain_grad().

Consider the small computation $x \rightarrow h = 3x \rightarrow o = \sum_i h_i^2$. Predict the attributes of the three tensors before running the cell. The exact name printed for grad_fn does not matter; what matters is whether it is None.

References: leaf tensors, leaf and non-leaf gradients, autograd mechanics.

tensor_x = torch.tensor([1., 2.], requires_grad=True)
tensor_h = 3 * tensor_x
tensor_objective = tensor_h.square().sum()

for name, t in [("x", tensor_x), ("h", tensor_h), ("objective", tensor_objective)]:
    print(name, "requires_grad:", t.requires_grad, "is_leaf:", t.is_leaf,
          "grad_fn:", t.grad_fn, "retains_grad:", t.retains_grad)
print("x.grad before backward:", tensor_x.grad)
assert tensor_x.is_leaf and tensor_x.grad_fn is None
assert not tensor_h.is_leaf and tensor_h.grad_fn is not None
assert tensor_x.grad is None

x requires_grad: True is_leaf: True grad_fn: None retains_grad: False
h requires_grad: True is_leaf: False grad_fn: <MulBackward0 object at 0x11247bbe0> retains_grad: False
objective requires_grad: True is_leaf: False grad_fn: <SumBackward0 object at 0x1117c3040> retains_grad: False
x.grad before backward: None

So far nothing has been differentiated, and x.grad is still None. Calling backward() starts from the objective, with derivative 1, and applies the chain rule back to x: here $\partial o/\partial h = 2h = [6, 12]$ and $\partial o/\partial x = 3 \cdot 2h = [18, 36]$.

To see the gradient with respect to h as well, call h.retain_grad() before backward(). This only lets us look at it; the gradient that reaches x is the same either way.

tensor_h.retain_grad()
print("Before backward:", tensor_x.grad, tensor_h.grad)
tensor_objective.backward()
print("After backward, x.grad:", tensor_x.grad)
print("After backward, h.grad:", tensor_h.grad)
print("h is still a non-leaf:", not tensor_h.is_leaf)
assert tensor_h.retains_grad and not tensor_x.retains_grad
torch.testing.assert_close(tensor_h.grad, torch.tensor([6., 12.]))
torch.testing.assert_close(tensor_x.grad, torch.tensor([18., 36.]))

Before backward: None None
After backward, x.grad: tensor([18., 36.])
After backward, h.grad: tensor([ 6., 12.])
h is still a non-leaf: True

A None gradient is not a zero gradient. It means that backward() never reached the tensor, or that the tensor is an intermediate result whose gradient was not retained. A leaf that is reached always gets a tensor, possibly full of zeros.

The stored gradients and the recorded computation are separate things. Setting x.grad = None clears the gradient and leaves the record alone. Each forward pass records new operations, and backward() normally frees what it needed from them (the tensor still shows a grad_fn afterwards).

Do not confuse h.retain_grad() with objective.backward(retain_graph=True). The first stores an intermediate gradient. The second keeps the record alive so that we can call backward() through it a second time, at the cost of memory; it does not clear or retain any .grad. In normal training we never need it, because each step runs a fresh forward pass. See backward.

Cutting a connection

Three operations that look similar but do different things:

• clone() copies the data, and the copy stays connected: gradients flow back through it to the original.
• detach() returns a tensor with no connection to the recorded computation. It does not copy anything: the two tensors share memory. Use detach().clone() for an independent snapshot.
• Inside a torch.no_grad() block, new operations are not recorded. Tensors that already exist are unchanged.

Calling requires_grad_() on a detached tensor makes it a new leaf; it does not restore the old connection. See detach.

source = torch.tensor([1., 2.], requires_grad=True)
connected_copy = source.clone()
detached = source.detach()
with torch.no_grad():
    untracked = 3 * source

for name, t in [("source", source), ("clone", connected_copy),
                ("detach", detached), ("no_grad result", untracked)]:
    print(name, "requires_grad:", t.requires_grad, "is_leaf:", t.is_leaf,
          "grad_fn:", t.grad_fn)
assert source.requires_grad
assert connected_copy.grad_fn is not None
assert detached.grad_fn is None and not detached.requires_grad
assert untracked.grad_fn is None and not untracked.requires_grad
connected_copy.sum().backward()
torch.testing.assert_close(source.grad, torch.ones(2))

source requires_grad: True is_leaf: True grad_fn: None
clone requires_grad: True is_leaf: False grad_fn: <CloneBackward0 object at 0x111fd7b20>
detach requires_grad: False is_leaf: True grad_fn: None
no_grad result requires_grad: False is_leaf: True grad_fn: None

Optional: autograd detective

Four short cases that use the previous section. For each one, predict what happens before running it: an error, a missing (None) gradient, or a number. Then decide whether that is what the author wanted. Not every case is a bug.

If a fix is needed, make the smallest one and add an assertion on the gradient you expect. Each case creates its own tensors, so you can run it on its own. The last case catches its error, so that running the whole notebook does not stop there.

1. Did the gradient stop at the intermediate tensor?

We want the gradient of sum(hidden**2) with respect to both leaf and hidden. After backward(), what are grad_fn, is_leaf and .grad for each of them? Does hidden.grad is None mean that the gradient did not flow through hidden? Rebuild the computation so that both gradients are available. PyTorch prints a warning when we read hidden.grad; the warning is part of the answer.

leaf = torch.tensor([1., 2.], requires_grad=True)
hidden = 3 * leaf
hidden.square().sum().backward()
print("leaf:", leaf.is_leaf, leaf.grad_fn, leaf.grad)
print("hidden:", hidden.is_leaf, hidden.grad_fn, hidden.grad)

# Rebuild, retain the missing gradient, and check both values.

leaf: True None tensor([18., 36.])
hidden: False <MulBackward0 object at 0x110b46260> None

/var/folders/w9/wt0x5nzx3jngfvpdt_2j15p00000gn/T/ipykernel_43124/1142738744.py:5: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more information. (Triggered internally at /Users/runner/work/pytorch/pytorch/build/aten/src/ATen/core/TensorBody.h:493.)
  print("hidden:", hidden.is_leaf, hidden.grad_fn, hidden.grad)

2. Can we reconnect a detached tensor?

We want the gradient with respect to original. Someone detached an intermediate result, then turned gradients back on. Will backward() fail? Which tensor receives a gradient? Fix the code so that the gradient reaches original.

original = torch.tensor([1., 2.], requires_grad=True)
restart = (3 * original).detach().requires_grad_()
restart.square().sum().backward()
print("restart:", restart.is_leaf, restart.grad_fn, restart.grad)
print("original.grad:", original.grad)

# Restore the intended differentiable path and check original.grad.

restart: True None tensor([ 6., 12.])
original.grad: None

3. Is accumulation always a mistake?

We want the gradient of the mean squared error over all three examples, but we process them in two microbatches of different sizes. The parameter does not change between the two backward() calls.

Predict parameter.grad after each call. Should we clear it between the microbatches? Compare with a single backward() on the full batch. Then replace the normalization with errors.square().mean() / len(microbatches): why does the result change here?

parameter = torch.tensor(1., requires_grad=True)
microbatches = [torch.tensor([1., 2.]), torch.tensor([3.])]
num_examples = sum(len(values) for values in microbatches)
for values in microbatches:
    errors = parameter * values  # All targets are zero.
    (errors.square().sum() / num_examples).backward()
    print("Accumulated gradient:", parameter.grad)

# Compare with the gradient of a single full-batch mean loss.

Accumulated gradient: tensor(3.333)
Accumulated gradient: tensor(9.333)

4. Is every in-place change safe inside no_grad?

To differentiate a square, autograd needs the original value of mutable. What happens if we change that value before calling backward(), even inside no_grad? Reorder the code so that it takes one gradient step, with learning rate 0.1, from the original value.

Then try modifying mutable.detach() in place before backward(). Does detaching protect the original data?

mutable = torch.tensor([1., 2.], requires_grad=True)
stale_loss = mutable.square().sum()
with torch.no_grad():
    mutable.add_(1.)
try:
    stale_loss.backward()
except RuntimeError as error:
    print("Backward raised RuntimeError:", str(error).splitlines()[0])

# Rebuild from [1., 2.], differentiate, then update under no_grad.

Backward raised RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [2]] is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True, check_nan=False).

Autograd detective: diagnoses and repairs

1. The gradient did pass through hidden. leaf is a leaf, with grad_fn=None and gradient [18, 36]. hidden is not a leaf and has a backward node. Its gradient [6, 12] was computed along the way, but not stored. Call retain_grad() before backward, on a fresh computation.

leaf = torch.tensor([1., 2.], requires_grad=True)
hidden = 3 * leaf
hidden.retain_grad()
hidden.square().sum().backward()
assert leaf.is_leaf and not hidden.is_leaf
torch.testing.assert_close(leaf.grad, torch.tensor([18., 36.]))
torch.testing.assert_close(hidden.grad, torch.tensor([6., 12.]))

2. Backward succeeds, on the wrong leaf. restart.grad is [6, 12], while original.grad stays None. Turning gradients back on after detach() does not restore the operation that connected the two tensors. Keep that operation in the recorded computation.

original = torch.tensor([1., 2.], requires_grad=True)
connected = 3 * original
connected.square().sum().backward()
torch.testing.assert_close(original.grad, torch.tensor([18., 36.]))

3. Here accumulation is intended. The gradient is 10/3 after the first call and 28/3 after the second. Each microbatch contributes its sum of squared errors divided by the total number of examples, so clearing between the calls would throw away the first contribution. Dividing each microbatch mean by two instead gives 11.5: the microbatch with one example then counts as much as the one with two.

full_parameter = torch.tensor(1., requires_grad=True)
all_values = torch.cat(microbatches)
(full_parameter * all_values).square().mean().backward()
torch.testing.assert_close(parameter.grad, full_parameter.grad)
torch.testing.assert_close(parameter.grad, torch.tensor(28. / 3.))

4. Backward raises an error. The in-place update changed a value that backward needs. Autograd detects this; no_grad only means the update itself is not recorded. Move the update after backward(). Modifying mutable.detach() in place would change the same memory and fail in the same way.

mutable = torch.tensor([1., 2.], requires_grad=True)
mutable.square().sum().backward()
torch.testing.assert_close(mutable.grad, torch.tensor([2., 4.]))
with torch.no_grad():
    mutable -= 0.1 * mutable.grad
mutable.grad = None
torch.testing.assert_close(mutable, torch.tensor([0.8, 1.6]))
assert mutable.is_leaf and mutable.requires_grad

Optional: views, copies and in-place operations

Some indexing operations return a view, which shares memory with the original tensor; others return a copy. Slicing gives a view, while indexing with a list gives a copy:

A = torch.arange(6.).reshape(3, 2)
view = A[:2]
copy = A[[0, 1]]
view[0, 0] = -10
print("Original:\n", A)
print("Copy made before the update:\n", copy)
assert A[0, 0].item() == -10 and copy[0, 0].item() == 0

# An underscore usually marks an in-place operation.
A.add_(1)
print(A)

Original:
 tensor([[-10.,   1.],
        [  2.,   3.],
        [  4.,   5.]])
Copy made before the update:
 tensor([[0., 1.],
        [2., 3.]])
tensor([[-9.,  2.],
        [ 3.,  4.],
        [ 5.,  6.]])

In-place operations need care when autograd is involved: backward() may need a value that the operation overwrites (see case 4 of the autograd detective). The update in Section 4 is safe because it runs after backward().

A tensor is stored as a flat block of memory plus strides: how many elements to skip to move one step along each axis. A transpose swaps the strides without moving any data, so the result is no longer contiguous. reshape returns a view when it can and copies otherwise. We will come back to memory layout when we talk about performance.

A = torch.arange(6.).reshape(3, 2)
print("A:", A.shape, A.stride(), A.is_contiguous())
print("A.T:", A.T.shape, A.T.stride(), A.T.is_contiguous())
print("Flattened transpose:", A.T.reshape(-1))

A: torch.Size([3, 2]) (2, 1) True
A.T: torch.Size([2, 3]) (1, 2) False
Flattened transpose: tensor([0., 2., 4., 1., 3., 5.])

Optional: a few array operations

Three short functions to write yourself. Each check cell tells you whether your version is still missing; write the function without loops over the elements, and pay attention to the edge cases the checks include.

First, sum the main diagonal and the anti-diagonal of a square matrix. When the size is odd, the center element is on both diagonals, so it is counted twice.

def sum_diagonals(A: Float[torch.Tensor, "n n"]) -> Optional[Float[torch.Tensor, ""]]:
    # Replace None with your implementation.
    return None

def sum_diagonals(A: Float[torch.Tensor, "n n"]) -> Optional[Float[torch.Tensor, ""]]:
    return A.diagonal().sum() + A.flip(dims=[1]).diagonal().sum()

if sum_diagonals(torch.eye(2)) is None:
    print("Write sum_diagonals first.")
else:
    torch.testing.assert_close(sum_diagonals(torch.arange(1., 10.).reshape(3, 3)), torch.tensor(30.))
    torch.testing.assert_close(sum_diagonals(torch.arange(1., 17.).reshape(4, 4)), torch.tensor(68.))
    torch.testing.assert_close(sum_diagonals(torch.tensor([[5.]])), torch.tensor(10.))
    # The two diagonals of the matrices above have equal sums; this one does not.
    torch.testing.assert_close(sum_diagonals(torch.tensor([[1., 2.], [3., 5.]])), torch.tensor(11.))
    print("sum_diagonals passed: odd, even, 1x1 and asymmetric matrices.")

sum_diagonals passed: odd, even, 1x1 and asymmetric matrices.

Next, normalize each row to mean zero and variance one. Variance here is the average squared deviation (correction=0). A constant row cannot be rescaled to unit variance, so it must become all zeros, without producing nan.

def normalize_rows(A: Float[torch.Tensor, "batch features"], eps: float = 1e-8) -> Optional[Float[torch.Tensor, "batch features"]]:
    # Replace None with your implementation.
    return None

def normalize_rows(A: Float[torch.Tensor, "batch features"], eps: float = 1e-8) -> Optional[Float[torch.Tensor, "batch features"]]:
    centered = A - A.mean(dim=1, keepdim=True)
    std = A.std(dim=1, correction=0, keepdim=True)
    return centered / std.clamp_min(eps)

A = torch.tensor([[1., 2., 4.], [4., 4., 4.]])
if normalize_rows(A) is None:
    print("Write normalize_rows first.")
else:
    normalized = normalize_rows(A)
    print(normalized)
    torch.testing.assert_close(normalized.mean(dim=1), torch.zeros(2), atol=1e-6, rtol=0)
    torch.testing.assert_close(normalized.var(dim=1, correction=0), torch.tensor([1., 0.]))
    # A square input catches normalization over the wrong axis.
    square = torch.tensor([[1., 3.], [2., 10.]])
    torch.testing.assert_close(normalize_rows(square), torch.tensor([[-1., 1.], [-1., 1.]]))
    print("normalize_rows passed, including the constant row and a square input.")

tensor([[-1.069, -0.267,  1.336],
        [ 0.000,  0.000,  0.000]])
normalize_rows passed, including the constant row and a square input.

Finally, compute the cosine similarity between every pair of rows. Hint: normalize each row to unit length first; then a single matrix product gives all the pairs. We will meet this operation again with attention layers.

Cosine similarity is undefined for a zero vector. Following F.normalize, we leave zero vectors at zero, so their similarity with every row is zero.

def cosine_similarity(A: Float[torch.Tensor, "batch features"]) -> Optional[Float[torch.Tensor, "batch batch"]]:
    # Replace None with your implementation.
    return None

def cosine_similarity(A: Float[torch.Tensor, "batch features"]) -> Optional[Float[torch.Tensor, "batch batch"]]:
    unit = F.normalize(A, p=2, dim=1, eps=1e-12)
    return unit @ unit.T

A = torch.tensor([[1., 0.], [0., 1.], [-1., 0.], [0., 0.]])
if cosine_similarity(A) is None:
    print("Write cosine_similarity first.")
else:
    expected = torch.tensor([[1., 0., -1., 0.], [0., 1., 0., 0.],
                             [-1., 0., 1., 0.], [0., 0., 0., 0.]])
    torch.testing.assert_close(cosine_similarity(A), expected)
    # Scaling a row must not change its similarities; the shape is (batch, batch).
    B = torch.tensor([[3., 4.], [6., 8.], [4., -3.]])
    torch.testing.assert_close(cosine_similarity(B), torch.tensor([[1., 1., 0.], [1., 1., 0.], [0., 0., 1.]]))
    assert cosine_similarity(torch.randn(5, 3)).shape == (5, 5)
    print(cosine_similarity(A))

tensor([[ 1.,  0., -1.,  0.],
        [ 0.,  1.,  0.,  0.],
        [-1.,  0.,  1.,  0.],
        [ 0.,  0.,  0.,  0.]])

Further reading

• PyTorch tensor tutorial
• Autograd mechanics
• Tensor views