PT01: Introduction to PyTorch

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 student version: answers to the check and to the optional exercises are not included.

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

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__)

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)

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)

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)

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)

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)

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.

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

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

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)

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.

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.

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.

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.

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

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.]))

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))

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.

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.

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.

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.

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)

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))

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

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.")

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

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.")

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

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))

Further reading

• PyTorch tensor tutorial
• Autograd mechanics
• Tensor views