Reference sheet

PyTorch calls you will use

The commands below appear in this project or are useful for checking it. Here N is config.modulus, d_vocab = N + 1, and model inputs have shape (batch, 3, N + 1).

First, read the shapes

A tensor's shape lists its axis lengths. For x.shape == (B, 3, N + 1), there are B equations, three token positions per equation, and N + 1 vocabulary entries per token. Axis 0 is batch, 1 is position, and 2 or -1 is vocabulary. A negative axis counts from the end: -2 is position here.

dim tells PyTorch which axis an operation uses. In scores.softmax(dim=-1), each row is normalized across its last axis. In values.sum(dim=-1), that axis disappears unless keepdim=True.

Create and encode data

Call	Arguments to supply	Result in this project
`torch.tensor(data, dtype=...)`	`data` is a number or nested list; optional `dtype` selects the number type.	`torch.tensor([0, 1, N])` makes three integer token IDs when `N` is the configured modulus. Integer IDs are suitable for one-hot encoding.
`torch.arange(end)`	`end` is excluded.	`torch.arange(config.modulus)` gives IDs `0, …, N−1`. You can also pass `start, end, step`.
`torch.meshgrid(a, b, indexing="ij")`	`a` and `b` are 1D tensors; use `indexing="ij"` so axes follow input order.	Two `(N, N)` grids holding every ordered pair of numbers.
`torch.randperm(n)`	`n` is a nonnegative integer.	A shuffled ordering of `0, …, n−1`; use it to sample pairs without replacement.
`torch.manual_seed(seed)`	`seed` is an integer.	Sets PyTorch's random seed so a run can reproduce the same shuffle and initialization.
`torch.nn.functional.one_hot(ids, num_classes=config.d_vocab)`	`ids` must contain integer class IDs from `0` through `N`; `num_classes` fixes output width.	Appends an axis of length `N + 1`. A `(batch, 3)` input becomes `(batch, 3, N + 1)`. It returns integer values; use `.to(torch.float32)` for model math.

Shape and combine tensors

Call	Arguments to supply	What it does
`torch.stack((a, b, c), dim=-1)`	Same-shaped tensors and the axis where a new dimension goes.	Three `(N, N)` grids become one `(N, N, 3)` tensor of equations.
`t.reshape(N**2, 3)`	The desired dimensions; their product must equal the number of elements.	Flattens an `(N, N, 3)` grid into `N²` equations.
`t.unsqueeze(dim)`	`dim` is where to insert an axis of length 1.	`(3, N + 1)` with `unsqueeze(0)` becomes a one-example batch `(1, 3, N + 1)`.
`t.squeeze(dim)`	Optional `dim` selects a size-1 axis to remove.	Be careful: `squeeze()` with no `dim` also removes a batch axis when batch size is 1.
`t.mT`	No arguments.	Swaps the last two axes of a matrix or batch of matrices; used for `Q @ K.mT`.
`t.to(torch.float32)`	A target dtype; alternatively pass a device such as `"cpu"`.	Converts integer one-hot values to floating point for multiplication with weights.

Compute the forward pass

Call	Arguments to supply	What it does
`a @ b`	Two tensors whose matrix dimensions align.	Matrix multiplication. For example, `x @ W_E` maps vocabulary vectors to hidden vectors; `*` means elementwise multiplication instead.
`torch.relu(t)`	A numeric tensor.	Replaces negative values with zero; used in the MLP.
`t.max(dim=-1, keepdim=True).values`	The axis to reduce and whether to retain it at length 1.	Finds each row's maximum. Keeping the axis lets the maximum broadcast back across that row before softmax.
`t.softmax(dim=-1)`	The axis of competing choices.	Turns attention scores into nonnegative weights that sum to 1 across that axis.
`t.log_softmax(dim=-1)`	For logits, use the vocabulary axis `-1`.	Computes log probabilities directly; used in cross-entropy for better numerical stability than taking `log(softmax(...))`.
`t.sum(dim=...)` / `t.mean(dim=...)`	The axis to reduce; omit `dim` to reduce all elements.	Sum over heads or vocabulary; average per-example losses into one scalar.

PyTorch broadcasts size-1 axes during compatible arithmetic. Write down shapes before using a broadcast; a shape that runs is not automatically the shape you intended.

Parameters and optimization

Call	Arguments to supply	What it does
`torch.randn(shape)`	Dimensions, for example `torch.randn(config.d_vocab, config.d_model)`.	Creates random normal values to initialize a weight matrix.
`nn.Parameter(t)`	A tensor `t` containing initial values.	Registers a learnable tensor on an `nn.Module`, so `model.parameters()` finds it.
`torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=1.0)`	Learnable parameters, learning rate `lr`, and decoupled weight decay strength.	Creates the optimizer used by this project. The example values are from `config.py`, not universal defaults.
`optimizer.zero_grad()`	No required arguments.	Clears gradients left by the previous step; otherwise new gradients accumulate.
`loss.backward()`	Here `loss` is one scalar tensor.	Backpropagates through recorded operations and fills each parameter's `.grad`.
`optimizer.step()`	No required arguments for AdamW here.	Updates parameters using their gradients. Call it after `backward()`.
`with torch.no_grad():`	A Python block, not a function argument.	Stops recording gradients while measuring train or held-out performance.

optimizer.zero_grad()
logits = model(inputs)
loss = loss_fn(targets, logits)
loss.backward()
optimizer.step()

Useful checks in tests

Call	Arguments to supply	Question it answers
`t.shape` / `t.dtype`	No arguments.	Does the tensor have the promised dimensions and number type?
`t.argmax(dim=-1)`	The axis containing competing scores.	Which token has the largest score? Use the last axis to decode one-hot vectors or logits.
`torch.equal(a, b)`	Two tensors.	Are shape and all values exactly equal? Good for integer IDs.
`torch.isclose(a, b, atol=1e-6)`	Two numeric tensors and an optional absolute tolerance.	Are floating-point values close enough for a small calculation?
`torch.isfinite(t).all()`	A numeric tensor.	Are all results finite, with no `nan` or infinity?

Start with one hand-computable input, then test a boundary or shape case. The project task asks each group to add a new test beyond the reference tests.