PyTorch, Beyond the API: Shapes, Storage, Gradients, and State

type

status

date

slug

summary

1. Shape is part of the algorithm

Consider predictions that perfectly match their targets:

The first expression compares every prediction with every target. Broadcasting aligns trailing dimensions, so (3, 1) and (3,) produce (3, 3), not three matched pairs. The code runs. The objective is wrong. PyTorch's broadcasting rules specify this behavior.

The important question is not “Are the dimensions compatible?” It is “Does the resulting computation preserve the relationship I intended?”

For a language model, I annotate logits as (batch, time, vocabulary), not just (B, T, V). Then softmax(dim=-1) means choosing among vocabulary items for each position. Applying it over time is a different calculation, even though the output shape is unchanged.

Small shape habits help. squeeze(-1) removes the axis I named; an unrestricted squeeze() can also remove the batch axis when the batch size becomes one. keepdim=True preserves a reduced axis so its meaning remains visible. The squeeze documentation explicitly warns about accidentally removing a batch dimension. In PyTorch reductions, the documented keyword is keepdim, not NumPy's keepdims.

Indexing selects coordinates, not necessarily a rectangle

One question in my fundamentals notes was why two lists of indices did not return a submatrix:

In the first expression, the index arrays have the same shape. They specify three coordinate pairs: (0, 2), (1, 3), and (2, 4).

In the second, (3, 1) row indices broadcast with (3,) column indices. That constructs all nine combinations. “Pairwise versus Cartesian” is therefore a consequence of the index shapes, not a rule that list indexing always pairs elements.

There is also a storage distinction. Basic slicing normally returns a view; advanced indexing returns a copy when reading. But w[rows, cols] = value writes into w itself. Do not confuse reading an indexed result with indexed assignment. The tensor-view documentation explicitly distinguishes them.

`gather` becomes simple when written as an equation

For a two-dimensional input, gathering along dimension one means:

Only the selected coordinate changes. The other coordinate still comes from the output position. The API contract also makes two useful constraints explicit: the output has the index tensor's shape, and input and index do not broadcast against each other.

The repeated selection of 50 contributes twice to its gradient. Selecting values can be differentiable with respect to those values. That does not make the integer indices themselves differentiable.

In a language model, the same operation selects the probability of each observed token:

The inserted singleton axis is not a trick to memorize. It makes the index tensor say: “At every batch and time coordinate, select one vocabulary entry.”

2. A new tensor object does not necessarily mean new storage

A strided tensor can be understood as values in storage plus instructions for locating them: shape, strides, and an offset. Transposing can change those instructions without moving the values.

This explains the difference between view and reshape. view requires a layout compatible with the requested shape. reshape can use a view when possible and copy otherwise. A tensor does not have to be globally contiguous for every possible view to work; the actual condition concerns its sizes and strides.

For the transposed tensor above, y.view(-1) fails, while y.reshape(-1) produces [0, 3, 1, 4, 2, 5] using a copy. Adding contiguous() before view() explicitly permits that layout conversion.

My rule is to use reshape when I need the new shape and accept a possible copy. I use view when sharing the existing layout is part of the intended contract. Neither name proves that the surrounding computation is efficient; that requires measurement.

Copying and detaching are independent operations

This is one of the most useful distinctions in the entire library:

Operation	Shares input storage?	Retains a differentiation path to the input?
A normal tensor view	Yes	Yes, when grad recording applies
`x.clone()`	No	Yes, when grad recording applies
`x.detach()`	Yes	No
`x.detach().clone()`	No	No

Cloning is differentiable. Detaching breaks the differentiation relationship but shares storage.

A useful consequence follows: copying a tensor does not inherently stop learning, and stopping gradients does not inherently protect the original values from mutation. For an independent snapshot, use both operations.

The same question appears at the NumPy boundary. torch.from_numpy(a) shares supported CPU array storage. torch.tensor(a) copies its input. NumPy operations do not become part of PyTorch's autograd graph merely because the two objects share memory.

In-place operations are about old values, not aesthetics

A trailing underscore usually signals mutation, as in add_() or zero_(). Assignment into a tensor can mutate it too.

The danger is not that all mutation is incompatible with autograd. The danger is overwriting a value that backward still needs. PyTorch tracks saved tensors and their versions to detect many such cases. Initializing a parameter before its first forward pass and changing a saved intermediate after forward are very different situations. See autograd's in-place-operation discussion.

When debugging, I therefore ask “Who still needs the previous value?” before asking whether an underscore might save memory.

3. A gradient is not proof that the optimizer will update a tensor

My parameter-registration notes contained an important ambiguity: they treated registration and participation in backpropagation as if they were the same thing.

They are not. Separate three questions: Is there a differentiable path from the loss? Is the tensor registered on a module? Does the optimizer hold that exact object?

Both tensors receive gradients. Only one is included in the optimizer constructed from model.parameters().

nn.Parameter provides special registration behavior when assigned as a module attribute. A regular tensor with requires_grad=True does not acquire that behavior. Conversely, an optimizer can receive a suitable ordinary leaf tensor explicitly. Autograd does not require a tensor to belong to a model registry.

A registered parameter may also be unused in a particular forward pass, or frozen. Registration alone does not establish a live path to the current loss.

Containers are part of model structure

Use ModuleList or ModuleDict for child modules. Use ParameterList or ParameterDict for parameters. A plain Python list does not register its contents merely by being attached to a module.

Sequential adds another behavior: it defines an ordered chain of forward calls. ModuleList registers children but leaves execution to your own forward. That difference matters for residual connections, branches, and blocks that return a cache as well as an activation.

Buffers cover model state that should be tracked without being returned as a model parameter. Examples include a fixed mask or running statistics. Registered buffers follow module device movement; persistent ones also appear in state_dict. Ordinary tensor attributes do not automatically get those services. These contracts are documented on nn.Module.

Why does `model.parameters()` return a generator?

I asked this separately, but it belongs to the same story. The API exposes an iterator over registered parameter objects. Its implementation yields them rather than first constructing a reference list.

This is not a mechanism for streaming model weights out of storage. list(model.parameters()) normally creates a list of references, not a second copy of all the weights. The avoided allocation is the reference container, which scales with the number of parameter objects, not their total number of scalar elements.

The module implementation also makes parameter deduplication visible. That matters for tied weights: two names can refer to one parameter object.

Optimizers remember objects, not attribute names

Suppose an optimizer was constructed and then model.weight was replaced with a new nn.Parameter. The module now knows the new object. The optimizer may still know only the old one. A gradient on the new object will not fix that mismatch.

Likewise, registering a parameter after optimizer construction does not automatically add it to the optimizer. add_param_group is an explicit mechanism for adding parameters; replacing or rebuilding parameter groups also requires thinking about existing optimizer state. The optimizer implementation shows the concrete object lists it maintains.

For initialization, prefer filling the existing parameter, for example nn.init.normal_(layer.weight, std=0.02). Assigning an ordinary tensor to an already registered parameter attribute generally raises a TypeError; it is not simply a silent way to “lose registration.” Creating a fresh parameter is allowed, but changes identity.

To read nn.Module from scratch, follow one small model through initialization, attribute assignment, calling the model, parameter traversal, and saving state. The useful source landmarks are __init__, __setattr__, the call machinery, named_parameters, and state_dict. Call model(x) in normal use; directly calling forward bypasses the module's call machinery, including relevant hooks.

4. Read a transformer as a sequence of shape contracts

The nanoGPT session review is where the earlier distinctions become one implementation.

Start with x of shape (B, T, C). nn.Linear(C, K) stores a weight of shape (K, C) and applies x @ weight.T + bias. It operates on the last axis, preserving the leading axes. Each output row is one learned combination of the input features.

That explains packed QKV:

The 3C dimension stores three different projections together. It does not permanently triple the transformer's hidden width.

For ordinary multi-head attention with H heads and D = C // H:

With fixed C, increasing H makes each head narrower. The reshape creates no extra information.

But “split the embedding into heads” needs one qualification. The split happens after learned projection. A head's projected coordinates can depend on all the input channels; heads are not restricted to disjoint chunks of the raw token embedding.

For a batched tensor, .T is not a synonym for “transpose the matrix inside each batch.” Name the two axes with transpose(-2, -1), or use a suitable matrix-transpose operation. Shape labels make this mistake visible before execution.

After attention, transpose the output from (B, H, T, D) back to (B, T, H, D) and reshape to (B, T, C). The output projection then mixes the concatenated head features.

Two temporary expansions, two different jobs

The feed-forward sublayer often follows C → 4C → C, with a nonlinearity in the middle. Unlike the 3C QKV packing, this is a hidden expansion for nonlinear feature transformation. The factor four is an architectural choice, not a PyTorch requirement.

A pre-normalized residual block has the form x + attention(norm(x)), followed by another residual update for the feed-forward sublayer. Moving normalization outside the addition changes the function and its gradient path. It is not an interchangeable rearrangement. The original build-nanogpt implementation is a useful concrete reference for these choices.

For LayerNorm(C), the default affine parameters are one scale and one bias per feature: 2C parameters, not one pair per batch or time position. Disabling affine parameters, or disabling the bias, changes that count. The same normalization rule is applied at each position, even though the position's hidden state may already contain information from other tokens.

Weight tying connects lookup and prediction

Let an embedding table E have shape (V, C). Input embedding selects rows: E[token_ids]. Output prediction scores vocabulary entries using hidden @ E.T.

The shapes fit because a linear layer mapping C → V also stores a (V, C) weight. Tying means assigning the same parameter object to both roles, not copying equal values into two independent objects. Gradients from the lookup path and the prediction path accumulate into that shared parameter.

This brings the four views together: compatible shapes, shared storage and identity, multiple gradient paths, and one parameter for the optimizer to update. The Embedding contract and the weight-tying assignment in the reference implementation make each part inspectable.

5. Backward is another computation worth understanding

Autograd removes the need to manually implement most derivatives. It does not remove the need to know what derivative the program should compute.

For real matrices Y = A @ B, let G be the upstream gradient with respect to Y. Then:

The shapes are a useful check, but not a proof. The derivation starts from dY = dA B + A dB. Substitute that into the Frobenius inner product dL = <G, dY> and regroup the terms multiplying dA and dB.

This extends directly to Y = A @ B @ C: the gradient of the middle matrix is A.T @ G @ C.T. For broadcasted batches, the backward calculation must also sum contributions over dimensions that were broadcast.

Why does `gradcheck` use numerical differences?

This was the most revealing question in my autograd discussions: why not have PyTorch calculate the derivative again?

Because a custom backward becomes the rule autograd uses. Asking the same route for another answer can repeat the same mistake. Finite differences check how the forward values change, without trusting the supplied backward rule.

For a scalar function, central differences estimate:

Here is a complete custom function to check:

Change the 3 to 2 and the check fails. That is the test doing something useful.

You can also give gradcheck an ordinary function composed of PyTorch operations. Its analytical derivative then comes from those built-in operations; writing a custom backward is not a prerequisite for using the checker. Gradcheck mechanics explains the analytical and numerical routes.

The numerical result is an approximation, not an oracle. The API notes recommend double precision for the default tolerances and warn about non-differentiable points and overlapping-memory inputs. Start with small deterministic inputs away from kinks. Passing a local gradient check does not prove that the chosen loss solves the intended problem.

Prefix products: direct gradients versus total gradients

Another note considered all prefix products:

Suppose a loss depends on every P[i], and G[i] is the direct upstream contribution at output i. The total gradient at P[i] also includes contributions through later prefixes.

Call that total A[i]. Work backward:

The subtle step is adding the direct contribution to the contribution flowing through later outputs. Forgetting it differentiates a different dependency graph.

This problem also connects to parallel scans. Matrix multiplication is associative, so prefix products can use a tree-like scan while preserving operand order. The backward recurrence can be represented as composition of affine maps X → X R + S. That composition is associative too.

A Hillis–Steele implementation uses logarithmically many dependent stages but more total work than the serial recurrence. It is a different work-versus-parallelism trade-off, not a free speedup. The companion code includes forward and backward scans and compares them with a serial implementation and autograd.

6. A KV cache changes coordinates, not the meaning of attention

Without a cache, a causal attention matrix for T tokens allows query i to see keys at positions at most i. With a cache, local query index zero may correspond to a much later absolute position.

Suppose three tokens are cached and the next call processes two new tokens. The query positions are 3 and 4; the key positions are 0 through 4. The correct allowed mask is:

A lower-triangular mask built without the offset is wrong. Removing the mask entirely is also wrong for the first query: it would see the later token in the new chunk.

This is why “there is a KV cache, so no causal mask is needed” is too broad. For one-token decoding, where the keys contain exactly the valid past tokens and the current token, every available key is allowed. Multi-token cached decoding still needs the correct causal restriction. Padding and other attention restrictions require additional treatment.

There is an API-specific trap as well. In scaled_dot_product_attention, a boolean True means allowed. That is the opposite of MultiheadAttention's boolean key-padding mask. Also pass dropout_p=0.0 for deterministic evaluation; the functional attention call applies the dropout probability it receives rather than independently consulting the model's evaluation mode.

In ordinary cached transformer decoding, each layer keeps its own K and V tensors. Append the new projections along the sequence axis, not the head axis. Positional embeddings or rotary positions must use the same absolute offset. A repeated torch.cat implementation is useful for learning correctness, but is not a claim about efficient production cache allocation.

Test equivalence, not just plausible text

A useful cache test compares full-sequence logits with the concatenated logits from cached calls. Test both one-token decoding and multi-token chunks. Put the model in evaluation mode, disable relevant randomness, and use numerical tolerances.

A second test changes future input tokens and checks that earlier logits do not change. This directly tests causality.

These checks are stronger than “the generated sentence looks reasonable.” They state properties that the implementation is supposed to preserve. The companion's tiny two-layer model passes both checks, including cached chunks of lengths 3, 1, 3.

7. The training loop connects the four views

Before tuning an optimizer, make the data-to-loss relationship explicit.

For a simple continuous token stream, take B*T + 1 tokens:

Position t predicts the following token. Shift once, either in data construction or inside the loss path, according to the model's contract. Shifting twice is as wrong as not shifting.

With logits (B, T, V) and class-index targets (B, T), a basic unpadded training step is:

cross_entropy expects logits, not probabilities from a preceding softmax. For explicit log-probabilities, log_softmax uses a formulation designed to avoid the numerical problems of separately computing softmax and then its logarithm.

backward() accumulates gradients. Resetting is therefore part of the algorithm, not cleanup. Intentional gradient accumulation is valid, but its normalization must match the intended effective batch. For unequal numbers of valid tokens, averaging microbatch means equally is not the same as averaging all valid tokens. zero_grad(set_to_none=True) also differs from filling every gradient with zero; optimizers can distinguish a missing gradient from a zero-valued one.

Masking chooses which examples count—and how much

Suppose sequence A has one valid token with loss 1, and sequence B has three valid tokens with loss 3 each.

The token-weighted mean is (1 + 3 + 3 + 3) / 4 = 2.5. The equally sequence-weighted mean is (1 + 3) / 2 = 2.0.

Neither is “just a different implementation of mean.” They encode different weighting choices. The denominator belongs in the algorithm's definition.

Decide what an all-masked batch should do. The companion's diagnostic helper raises an error. Other systems may explicitly skip such a batch, but silently producing a NaN is not a useful contract. Also avoid assuming that multiplying an invalid value by zero removes it: 0 * NaN is still NaN. Prevent invalid operations where possible rather than masking only their final outputs.

Initialization, precision, and evaluation mode

A positional parameter created with torch.empty has allocated storage but no meaningful initialization. Fill it before use. The normal initializer is nn.init.normal_, not nn.init.normalize_. These are small spelling and lifecycle issues with large consequences for a training run.

model.eval() changes the behavior of relevant modules, such as dropout. It does not disable autograd. torch.no_grad() controls graph recording; inference mode goes further and imposes restrictions on later autograd use. These are separate mechanisms, not interchangeable spellings of “inference.”

Mixed precision is another separate concern. Autocast selects suitable operation dtypes; loss scaling is a distinct mechanism, especially relevant to FP16 underflow. BF16's wider exponent range does not make every computation numerically safe. Gradient clipping cannot repair a forward pass that has already produced invalid values. For FP8, the exact scaling recipe and supported implementation matter; it is not simply a matter of casting every tensor to eight bits. PyTorch's AMP guide and numerical-accuracy notes are the starting references, rather than a universal low-precision recipe.

For debugging, inspect the first operation that produces a non-finite value. Separately check initialization, mask semantics, target alignment, the actual gradients, and optimizer membership. Then try to overfit a tiny fixed batch. A falling loss is useful evidence, but the earlier broadcasting example shows why it is not sufficient evidence of the right objective.

8. The same reasoning transfers beyond a transformer

Vectorization can reveal a simpler model

My nearest-neighbor notes asked how an L2 lookup could be expressed as a neural-network forward pass.

Expand squared distance:

For a fixed query, ||q||² does not affect which reference wins. So minimizing distance is equivalent, in exact arithmetic, to maximizing the affine score 2 q·x_i - ||x_i||².

The matrix formulation avoids explicitly building an (M, N, D) difference tensor, though the full score matrix still costs (M, N) storage. Chunking is necessary when that is too large. Floating-point cancellation and ties deserve separate tests.

Softmax preserves the winning exemplar in exact arithmetic. But averaging exemplar labels with softmax probabilities is not hard 1-nearest-neighbor classification. For probabilities [0.45, 0.30, 0.25] and labels [A, B, B], the nearest exemplar says A while aggregated class probability says B.

For L1 distance, the absolute value introduces another ingredient: |a| = ReLU(a) + ReLU(-a). The L2 affine trick does not transfer unchanged.

The useful skill is not deleting every Python loop. It is finding a representation whose computation and memory cost match the problem, then checking that the rewrite preserves its meaning.

Post-training makes the same distinctions more consequential

A GRPO-style training-loop note raised familiar issues in a new setting: batch, group, token, and vocabulary dimensions; gathered token log-probabilities; masks; detached rewards; and fixed old-policy probabilities.

Suppose rewards have shape (B, G), where G responses belong to each prompt. A within-prompt baseline reduces over the group axis, not across unrelated prompts. Standardization needs an explicit convention for variance, an epsilon, and a policy for degenerate groups.

For a policy ratio exp(new_log_prob - old_log_prob), the old log-probability is a fixed record of the sampling policy, not a quantity to recompute from the changing model and call “old.” A reward or advantage used as fixed policy-gradient feedback should not accidentally introduce a new differentiation path.

Now the earlier questions have direct consequences: Did gather select the sampled tokens? Are next-token positions aligned? Does the mask include prompt tokens? Are sequences or tokens weighted equally? Is the optimizer updating the intended policy parameters?

This is not a complete GRPO recipe. The source note contains a debugging prompt, not an original training implementation to repair. The transferable point is that tensor semantics are part of the learning objective. They are not a lower-level detail beneath it.

A few nanoGPT details that should remain implementation-specific

Some questions in the session review do not require another mental model, but they do require resisting overgeneralization.

Vocabulary padding. The GPT-2 vocabulary size in this implementation is 50,257; the training configuration uses a padded size of 50,304, a multiple of 64. The additional rows do not create tokenizer-defined tokens. A decoder should restrict sampling to valid IDs. Padding may help a particular matrix multiplication implementation, but the speedup needs measurement on the actual hardware. The GPT-2 encoder also explains why a Unicode character need not correspond to one token: bytes and vocabulary entries are different levels of representation.

Weight decay. Splitting parameters by dimensionality—decaying matrices while excluding common one-dimensional bias and normalization parameters—is a convention used in the reference code. It is not a theorem that one-dimensional tensors should never be regularized.

Token conversion. The loader's NumPy int32 conversion followed by a PyTorch long tensor is an implementation choice, not a universal requirement to copy token arrays twice. Embedding lookup supports integer index types including IntTensor and LongTensor; the usual class-index cross-entropy path expects long targets. Validate token range before any narrowing conversion.

Distributed data. A rank offset of B*T*rank, followed by advances of B*T*world_size, can assign disjoint input windows within a shard. Shard transitions and sequence boundaries still need coordinated handling. DistributedDataParallel synchronizes model training; it does not shard inputs for you.

NumPy habits. Check floating-point defaults, integer reductions, device placement, and type promotion when porting code. A CPU NumPy array and a CUDA tensor are not interchangeable. A negative-stride NumPy view may need a copy before conversion. Extracting a device scalar for Python control flow can force synchronization and complicate compilation. Seeding is useful, but does not promise identical results across every backend, precision, or release.

These details are reasons to read the implementation and its contracts, not reasons to memorize every incidental choice as a rule of deep learning.

How I would rebuild this understanding from scratch

I would keep a small notebook of claims and counterexamples, rather than a long list of function definitions.

First, write the intended calculation with a loop and named axes. Vectorize it and compare both the outputs and the gradients. Include batch size one, unequal sequence lengths, repeated indices, and an empty supervision mask.

Second, build the smallest possible model that separates gradient tracking, parameter registration, and optimizer membership. Intentionally leave one tensor unregistered. Then replace a parameter after creating the optimizer. Predict which values will change before running either example.

Third, implement one custom differentiable operation, check it numerically, and deliberately break its backward rule. Implement attention without a cache, then add one. Demand agreement between full and cached logits, not just a plausible sample.

Finally, read a compact real implementation again. The source should now answer concrete questions: where a parameter is registered, why an axis moves, which values backward needs, and what the optimizer actually owns.

The hardest confusion is treating shape, storage, gradients, and model state as if they were one thing. A copy can preserve gradients. A detached tensor can share memory. A tensor can receive a gradient without appearing in the optimizer. A correct-looking shape can express the wrong objective.

That is the standard I want from PyTorch fluency: not remembering every API, but being able to predict what a line changes—and design a small test that would prove me wrong.

Verification note

The companion pytorch_checks.py contains 30 passing CPU checks covering the central counterexamples, custom backward verification, matrix-product scans, gradient accumulation, weight tying, token alignment, masked reductions, and full-versus-cached transformer outputs. It was run during preparation of this draft with Python 3.13.5, PyTorch 2.10.0+cpu, and NumPy 2.3.5.

Those tests are not GPU benchmarks. CUDA performance, FP8/BF16 behavior, multi-process distributed training, and production-scale cache allocation were not tested. Some snippets in the article explain a local operation and assume the surrounding tensors; the companion file supplies runnable implementations and assertions.

Runnable appendix: complete source, tests, and a training demo

The two source files below are complete, not pseudocode. Save them in the same writable folder as pytorch_checks.py and train_demo.py. No datasets, model downloads, API keys, or GPU are required after installing the dependencies. The scripts run locally; the Notion code blocks display the source rather than execute it.

Install and run

The verification environment is Python 3.13.5, PyTorch 2.10.0+cpu, and NumPy 2.3.5 on Linux. The versions are pinned to this run, not presented as the latest releases. For platform-specific wheels, use the official PyTorch installation commands.

On Windows, activate with .venv\Scripts\Activate.ps1 in PowerShell. On a supported macOS Python installation, replace the two dependency-install lines with python -m pip install torch==2.10.0 numpy==2.3.5; the code still uses CPU. The reported run was Linux CPU; macOS and Windows were not tested here.

Run one check in isolation with standard Python unittest:

python pytorch_checks.py also writes validation.json. The demo writes demo_results.json. Use a writable working copy of the files. Do not run Python with -O when experimenting with assertion-based snippets.

What each group checks

Checks	What they verify
01–06	Broadcasting, pairwise indexing, repeated gather indices, views, copying, detachment, and NumPy sharing
07–13	Parameter registration, containers, iterator exhaustion, optimizer identity, linear layers, weight tying, and LayerNorm
14–16	Correct and deliberately wrong backward rules, serial matrix derivatives, and parallel scans
17–19	Full/cached logits, unchanged prefix logits when future tokens change, and causal masks with a cache offset
20–26	Cross-entropy values and gradients, masked loss weighting, nearest-neighbor equivalence, group advantages, token shifts, integer inputs, and sampling
27–30	Unsafe in-place mutation, gradient reset, evaluation versus no-grad, and a tiny-batch overfit check

A deliberate error is not an expected failing test. For example, test 14 succeeds only when gradcheck accepts Cube and rejects WrongCube; test 27 succeeds only when autograd detects unsafe mutation.

File 1: `pytorch_checks.py` — implementation and all 30 tests

Copy this entire code block into pytorch_checks.py. The first half defines the reusable tensor functions and tiny transformer; Checks contains the full tests. Importing the file does not run the test suite.

File 2: `train_demo.py` — train, evaluate, and generate

Copy this entire block into train_demo.py beside the first file. It imports TinyLM from pytorch_checks.py; no missing model classes or placeholder data need to be supplied.

The demo trains a two-layer character model on four shifted versions of the synthetic cycle abcd. It then generates 24 new tokens from abc, compares cached and full-prefix generation, and compares their logits. This is an intentional overfit smoke test, not a demonstration of natural-language understanding. TinyLM returns (logits, caches), so this executable training loop unpacks both values.

Results from the executed files

All 30 tests passed with zero failures and zero errors. They were rerun after retrieving the code package for this update.

The training demo produced this output in the verified CPU environment:

Loss decimals and runtime can vary with the platform. The meaningful checks are the asserted numerical agreement, the explicit failure cases, and the learned toy pattern. No CUDA speed, low-precision stability, distributed execution, or production-serving performance is established by these runs.

Try to break a claim

Replace 3 * x.square() with 2 * x.square() in Cube.backward: test 14 should fail. Remove the cache offset in CausalAttention.forward: the cached-equivalence tests should fail. Replace the parameter after constructing an optimizer and inspect test 10 before running it. These changes turn the appendix into a diagnostic exercise rather than a collection of examples that merely execute.

View	The question it answers
Shape and meaning	What does each axis represent? Which elements interact?
Storage	Where are the values? Which objects share those values?
Differentiation	Which operations connect the loss to a tensor?
Model and optimizer state	Which objects are registered, saved, moved, and updated?