Week 5 — Tensors and PyTorch: let the computer do the calculations

Tensor → store numbers and compute on groups of them

Autograd → apply the chain rule automatically

nn.Module → organize parameters and forward computation

Optimizer → update parameters

Focus on four things: tensors, shapes, Autograd, and the training loop.

For this course, think of a tensor as:

A multidimensional numerical container with shape, dtype, device, and related metadata.

Each part has a practical purpose:

Corresponding shapes:

scalar: []

vector: [3]

matrix: [2, 3]

A tensor is not a new set of arithmetic rules. It places Week 1's scalars, vectors, and matrices behind a common interface.

Why not just use Python lists?

Lists can store numbers, but a Python loop over millions of elements incurs interpreter overhead. Lists also do not provide a shared interface for tensor shapes, accelerator devices, and automatic differentiation.

A dense tensor gives numerical data a structured layout and passes whole operations to optimized lower-level kernels:

Python requests a matrix multiplication

A lower-level kernel performs many multiply-and-add calculations

Autograd records the dependency when gradient tracking is enabled and needed

Tensors are the common data representation linking the mathematics, accelerated operations, and backpropagation.

Suppose a batch contains four houses, each with three features:

dimension 0 = batch

dimension 1 = features

When working with text later, a common arrangement is:

[B, T, C]

When reading PyTorch code, trace shapes before interpreting the formulas.

Do not guess what an axis means

B, T, and C in [B,T,C] are project conventions, not names enforced by PyTorch. A tensor with shape [4,8,32] can be misinterpreted if one part of the code treats axis 1 as token position while another treats it as a channel. An operation may run despite the semantic mistake.

Add a shape trace beside complex code:

It serves as a readable annotation for dynamically shaped tensors.

Dtype

Neural-network parameters and activations commonly use floating-point numbers. This course uses integer torch.long token IDs for embedding lookup.

Why integer token IDs? They select rows of an embedding matrix, so they are indices. Why floating-point weights? Gradient-based training makes small numerical parameter changes that integer storage cannot represent.

Device

Tensors participating in an operation must satisfy its device requirements. A common mismatch is a model on GPU with its input still on CPU.

Moving data from CPU to GPU can require copying it into different memory. Repeated transfers have a cost. A common approach keeps the model on the chosen device, transfers each input/target batch there, and brings back only the values needed for reporting. The entire dataset need not fit on the accelerator.

1e-3 and math.exp(-3) are different: the first is 10⁻³; the second is the natural exponential e⁻³. Compare floating-point results with suitable tolerances, such as assert_close. Do not require bitwise equality across different orders of calculation.

Formula:

Suppose:

X: [4,3], four examples with three features each

W: [3,2], connecting three inputs to two neurons

b: [2]

One bias per neuron

Output:

Z: [4, 2]

Code:

@ performs matrix multiplication.

How much arithmetic is in this line?

X @ W applies two output neurons to each of four examples: 4×2=8 dot products. The bias is added to every row. A single line describes the batched forward calculation of the whole layer.

This is the value of the tensor interface: high-level code specifies a data transformation, while the kernel arranges the repeated arithmetic.

Z has shape [4,2], while b has shape [2].

PyTorch applies the same bias vector to every batch row:

[4, 2]

+ [2]

-------

[4, 2]

This is broadcasting.

Align shapes from the last axis. Aligned sizes must match or one must be 1; missing leading axes act as size 1 for this comparison.

Why broadcasting exists

The danger of broadcasting

Broadcastable does not mean semantically correct. For example, [B,T,C]+[T,1] can run even if the intended addition was a [C] bias. When a result is surprising, align shapes from the right and explain why each repeated axis is intended.

Return to Section 17 of Week 2: x=[1,2,3,4], y=[3,5,7,9], w=b=0, mean MSE, and η=0.01. Keeping these conditions should reproduce the same first step: loss=41, dw=−35, db=−12, and updated w=0.35, b=0.12. Tensors calculate the four items together; they do not change the learning rule.

Data:

The data follow:

Initialize parameters:

Forward：

The mathematics remains:

This means:

Track differentiable operations involving w so that we can later calculate the gradient of loss with respect to w.

In ordinary eager execution with gradient tracking enabled, PyTorch dynamically records the computational graph during forward.

w, b

↓

prediction

↓

error

↓

square

↓

mean

↓

loss

Which tensors need gradients?

For these examples, set requires_grad=True on the floating-point parameters to learn. Inputs and targets need not have their own gradients computed. Other tasks can require input gradients, so this is a choice for our training setup rather than a universal restriction.

requires_grad=True does not mean a gradient has already been calculated. It enables tracking for later differentiation. For a leaf parameter used by the loss, backward then accumulates the numerical result into .grad.

Leaf tensors and intermediate tensors

Directly created w and b used as parameters here are leaf tensors. Predictions, errors, and loss produced from them are intermediate tensors with grad_fn describing their differentiation history.

An optimizer normally holds references to persistent leaf parameters. Intermediate tensors connect the current computation and are not themselves the learned model state.

This line starts Week 4's backward calculation from the scalar loss.

Results are stored in:

These correspond to:

PyTorch does not guess the gradients. It follows the recorded dependencies and combines the operations' local derivatives using the chain rule.

Why usually call backward on a scalar loss?

The implicit starting seed for loss.backward() is:

For a multi-element output, backward needs a supplied gradient vector specifying how the output components contribute. A scalar mean loss defines one objective and allows an implicit seed of 1. For squared-error regression, average the per-example losses—not the signed errors that could cancel.

What .grad means

w.grad is neither the new weight nor a quantity to add directly to w. It is the current:

The optimizer uses it with the learning rate and any other update rules, such as momentum, to determine the actual parameter change.

Without an explicit gradient argument, loss.backward() normally requires a single-element output. If you have a vector of per-sample losses, first use .mean() or .sum() to define the objective. Vector outputs can also be differentiated, but the backward seed must be specified. Differentiation is not restricted to scalar-valued functions.

PyTorch accumulates gradients into .grad by default:

new grad = old grad + current grad

This is useful in some workflows, but a standard independent training step normally clears previous gradients:

When using an optimizer:

Otherwise, the next update includes contributions left over from previous backward calls.

Accumulation across backward calls supports deliberate schemes such as combining several small batches before one update. The training loop chooses when to clear that stored accumulation. This is distinct from summing multiple dependency paths within one backward calculation, which is mathematically required.

A clear standard sequence puts zero_grad() before forward. Clearing at the end of the previous completed update can also work, provided gradients are cleared at the intended boundary. Putting it at the start makes that boundary easy to see.

Before the first backward call, parameter.grad is usually None, so calling parameter.grad.zero_() unconditionally can fail. optimizer.zero_grad(set_to_none=True) clears stored gradients by setting them to None. No stored gradient and a stored all-zero tensor are different states. Week 11 discusses intentional accumulation.

Our update should not become another differentiable operation linking the training iterations:

The formula is unchanged:

Here w is a leaf tensor with requires_grad=True. An in-place update such as w -= ... in normal grad mode is rejected with a RuntimeError. Replacing it with the out-of-place w = w - ... can instead create a non-leaf tensor carrying history and connect iterations. Put the manual parameter update inside torch.no_grad(), or use the standard optimizer.step() shown here.

This implements the calculations from Weeks 2 and 4.

Larger models need a systematic way to manage parameters.

nn.Linear(1,1) already contains:

weight

bias

Call model(x), which invokes forward through the module's call machinery and preserves the expected hook behavior. Differentiable tensor operations inside forward are recorded when grad mode requires it.

Why nn.Module is more than a wrapper

A plain class can define forward too. nn.Module additionally maintains a registered model structure. Assigning an nn.Parameter or child module to self allows the framework to:

Collect registered parameters through model.parameters(); requires_grad determines whether each needs gradients

model.to(device)

Move registered parameters and buffers recursively between devices

Export parameters and persistent buffers through model.state_dict()

Set the training/evaluation mode of registered submodules

Registration lets the framework consistently manage the model's components.

nn.Linear(1,1) initializes parameters randomly, so its default first step will not match our hand calculation. For a fair Week 2 comparison, zero both weight and bias inside no_grad first. Its weight shape is [1,1]; the general storage rule is [outputs,inputs]. For row-wise samples, forward is X @ weight.T + bias. See course_examples/week05_three_ways.py for all three implementations.

Their responsibilities:

Model → produce predictions

Loss function → score predictions against targets

Autograd → calculate gradients

Optimizer → update parameters

The loss defines the objective; the optimizer defines how to use gradients to change parameters. A model can be trained with different objectives, and a loss can be paired with different optimizers. Backpropagation supplies the parameter gradients connecting them.

Now combine tensors, automatic differentiation, a module, and an optimizer. The following program is standalone. Because Linear normally initializes randomly, explicitly zero its parameters for comparison with Week 2.

The table's approximate values follow from the same fixed inputs. Printed floating-point values may differ in their final digits. In a loop, clear old gradients at the start of each independent step. zero_grad does not zero the learned parameters, and optimizer.step does not automatically clear .grad.

Running without an exception does not prove correct pairing. Predictions have shape [4,1]; targets accidentally stored as [4] can broadcast to [4,4], producing sixteen cross-comparisons instead of four matched pairs. Explicit shape checks are more informative than merely checking that loss is a scalar.

Next week, predictions are logits with shape [B,T,Vocab] and targets have shape [B,T]. The classification interface intentionally uses different shapes. Do not copy this regression equality check into it; each check must match the axis meanings and loss-function contract.

When matmul, reshape, or loss calculations fail, do not guess blindly.

Print at key boundaries:

Check each item:

In AI code, shape is part of the interface contract and data schema.

B, T, C, and V name axis sizes; they are not four new kinds of mathematics. First read [2,3,4] as “two sequences, three positions per sequence, four numbers per position,” then inspect which dimensions the next operation changes.

They change the behavior of layers such as Dropout and BatchNorm, rather than enabling or disabling Autograd.

A common validation pattern:

We will use this in more detail when training GPT.

Three controls that are easy to confuse:

model.train() → set registered layers to training behavior

model.eval() → set registered layers to evaluation behavior

torch.no_grad() → disable recording for backward within the context

eval() does not automatically disable gradient tracking, and no_grad() does not automatically turn off Dropout. Validation commonly uses both eval() and no_grad() because they address different requirements.

So far, the model receives numerical features.

GPT starts with text.

The next question:

How do we turn text into tensors, and how do those tensors support next-token prediction?

Week 6 builds this path:

Text

↓

Token IDs

↓

Embeddings

↓

Logits

↓

Cross Entropy Loss