JAX Pathways
Training Systems · Local CPULaptop CPULesson 3 of 4 in Phase 07

Phase 07 / Transformers & language models

A Transformer block and mixed precision

Revision v0.7.0 ·

Lesson · about 90 minutes · CPU practice

Build a residual attention block and compare numerical behavior across precision choices.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Masking, tokenization, and sequence packing
Hardware: CPU

Download notebook
Other downloadsPython worked examplesLesson text
Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder
py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment
python phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py

PowerShell command · repository folder
.\.venv\Scripts\python.exe phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Check: Output shape (4,8)(4, 8); finite values, zero-projection identity, compiled/eager agreement.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 07.03 Stage 1 of 5 · Understand
1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence
Stage 1 of 5 · Understand the MechanismCore mental model, equations & visual walkthrough

What you will be able to do

  • Follow the two residual paths
  • Normalize features rather than examples
  • Test causal independence and differentiation
  • Separate storage precision from sensitive arithmetic

The problem

An attention head mixes information across positions. What else makes it a Transformer block? We’ll add normalization, a small feature-processing network and two residual connections. Then we’ll inspect gradients and compare a lower-precision version. This is a deliberately small CPU teaching block, not a production architecture.

The idea

A Transformer block adds learned corrections to token representations. This lesson uses a single attention head, normalization before each branch, and two residual additions. Follow the implemented dimensions rather than assuming all Transformer architectures are interchangeable.

Follow the residual stream through the actual block

Four token positions each carry eight features. Query/key scores therefore form a four-by-four matrix. Weighted values and the output projection return to eight features, allowing the attention correction to be added to the input.

The feed-forward branch expands eight features to sixteen, applies GELU and returns to eight. It processes each token's feature vector; attention is the part that mixes eligible positions. Each residual addition requires matching shapes.

The saved output-minus-input heatmap shows the combined correction. Its equal rows arise from this fixture's constant-offset input and normalization behavior. They are not a universal property of attention. Use the architecture to identify which input change would test that explanation.

The actual single-head pre-LN block

Predict: What should happen when every learned projection is zero?

The center path carries four positions with eight features. Attention mixes eligible positions through four-by-four weights. The feed-forward branch changes each token from eight to sixteen features and back. Curved bypasses carry residual values to additions. No multihead or cache component is implied.
Architecture and dataflow mechanism diagram.

The center path carries four positions with eight features. Attention mixes eligible positions through four-by-four weights. The feed-forward branch changes each token from eight to sixteen features and back. Curved bypasses carry residual values to additions. No multihead or cache component is implied.

Open full-size diagram
LN⁡(x)i=γixi−μσ2+ϵ+βi,h(ℓ+1)=h(ℓ)+fℓ(LN⁡(h(ℓ)))\operatorname{LN}(x)_i = \gamma_i \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_i, \qquad h^{(\ell+1)} = h^{(\ell)} + f_\ell(\operatorname{LN}(h^{(\ell)}))

Follow the two residual paths

The first sublayer normalizes each token’s features, creates Q/K/V projections, performs causal attention and projects back to the model width. Add that result to the original xx. The second sublayer normalizes the intermediate residual, expands feature width, applies GELU, contracts back and adds again.

Every residual addition needs matching shapes. Zeroing all six projection matrices makes both sublayers contribute zero, so the block must return xx exactly. This is an independent boundary case that catches many incorrect residual connections.

x → norm → Q/K/V → causal attention → output projection
└────────────────────────────── add → residual
residual → norm → expansion → GELU → contraction
└────────────────────────────────── add → output

Normalize features rather than examples

For each token, compute the mean and variance over its feature axis. Subtracting the mean and dividing by sqrt(variance+epsilon) makes this sublayer insensitive to a shared finite feature shift, within numerical tolerance. Epsilon prevents division by zero for a constant feature row.

Layer normalization here does not compute statistics across the batch or sequence. Its result has the same shape as its input, and it has no running statistics. Batch normalization and learned affine normalization introduce different state or parameters; do not swap them without revisiting the contract.

Test causal independence and differentiation

Perturb the last token heavily and verify that the first three output rows stay the same. Pre-normalization and the MLP act per token, while the attention mask prevents earlier queries from reading the last key. A full gradient with respect to the parameter pytree should retain each leaf’s shape and remain finite for this input.

Finite gradients are necessary but not sufficient evidence of a correct block. Check known cases, changed inputs, residual identity and eager/compiled agreement together. The loss reduction determines the gradient scale.

Separate storage precision from sensitive arithmetic

Casting every parameter and input to bfloat16 is a simple controlled low-precision experiment, not a complete mixed-precision training policy. A common policy stores selected tensors at lower precision while computing reductions and normalization in float32. The second experiment compares fully cast bfloat16 computation with a policy that keeps feature normalization and score softmax in float32, while casting their outputs back for the lower-precision projections.

Record maximum absolute error for the supplied small input; do not require bitwise agreement. Precision affects both small denominators and accumulated dot products. If a comparison fails, inspect normalization and input scale before increasing the tolerance. Real models require workload-specific accuracy and performance qualification.

A residual block adds a correction to its input

Predict before looking: Does a residual connection force the output to equal the input?

A residual block adds a correction to its input. Rows are token positions and columns are features. Every cell shows output minus input at that location, so this is a map of the block’s correction, not its raw output. Purple indicates a positive correction, brown a negative correction, and near-white little change. Feature \(3\) increases by about \(0.0897\), while feature \(7\) decreases by about \(0.385\). Those same column patterns repeat across all four rows, up to floating-point rounding. The strongest color marks the largest change in magnitude, not the feature with the largest original activation.
Plot from a recorded CPU computation · JAX 0.9.2 · 1 CPU device

Read the figure

Rows are token positions and columns are features. Every cell shows output minus input at that location, so this is a map of the block’s correction, not its raw output. Purple indicates a positive correction, brown a negative correction, and near-white little change.

Feature 33 increases by about 0.08970.0897, while feature 77 decreases by about 0.3850.385. Those same column patterns repeat across all four rows, up to floating-point rounding. The strongest color marks the largest change in magnitude, not the feature with the largest original activation.

Connect it to the computation

Why do the rows look identical? The example builds input rows that differ only by a constant offset across their features. Per-token layer normalization removes that offset, so the normalized rows match. Their projected values match too, so attention’s weighted average is the same even when the causal mask allows different numbers of tokens. The subsequent normalized feed-forward branch preserves this symmetry.

The block therefore adds almost the same correction to different input rows; the final rows themselves are still different. This repeated pattern comes from the chosen input, not a rule that transformers treat every token identically. Change one feature of one token, rather than shifting the entire row, to break the symmetry and inspect how the correction changes.

Open full-size figureDownload data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Stage 2 of 5 · Build & Run the ReferenceStep-by-step code construction & verified execution output

Where to put & run this code: Work in phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py (or downloaded exercises/transformers-03.py / exercises/transformers-03.ipynb) and execute from the repository root with python3 phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py.

1 Prepare the experiment

In the CPU environment from setup, create main.py and add this block. Write down the dimensions and the known result before continuing.

main.py · append build step 1
# Step 1 — 1. Prepare the experiment: These imports provide JAX array transformations and NumPy...
import jax
import jax.numpy as jnp
import numpy as np

These imports provide JAX array transformations and NumPy reference calculations. Define the numerical functions next; the final build block supplies the concrete fixture and checks.

2 Implement the mechanism

Append this block in the same file. Follow the shape diagram above and identify each reduction or state transition.

main.py · append build step 2
# Step 2 — 2. Implement the mechanism: x → norm → Q/K/V → causal attention → output projection...
def layer_norm(x):
    # Reduce along axis=-1 to compute `mean`.
    mean = jnp.mean(x, axis=-1, keepdims=True)
    # Reduce along axis=-1 to compute `variance`.
    variance = jnp.mean((x-mean)**2, axis=-1, keepdims=True)
    return (x-mean) / jnp.sqrt(variance + 1e-5)

def init_block(key, width=8):
    # Create or split explicit PRNG key(s) (`keys`) for reproducible randomness.
    keys = jax.random.split(key, 6)
    shapes = [(width,width)]*4 + [(width,2*width),(2*width,width)]
    return {name: jax.random.normal(k,s)*0.1 for name,k,s in zip(
        ["q","k","v","o","up","down"], keys, shapes)}

def block_forward(p, x):
    h = layer_norm(x)
    q, k, v = h@p["q"], h@p["k"], h@p["v"]
    scores = q@k.T / jnp.sqrt(x.shape[-1])
    # Create evenly spaced index values in `allowed`.
    allowed = jnp.arange(x.shape[0])[:,None] >= jnp.arange(x.shape[0])[None,:]
    # Apply nonlinear activation or probability normalization to compute `weights`.
    weights = jax.nn.softmax(jnp.where(allowed, scores, -jnp.inf), axis=-1)
    residual = x + (weights@v)@p["o"]
    return residual + jax.nn.gelu(layer_norm(residual)@p["up"])@p["down"]

Code building blocks in this step

  • jnp.arange(n, dtype=...) — Creates a 1-D JAX array of evenly spaced values [0, 1, ..., n-1] on the target device.
  • x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
  • jax.random.PRNGKey(seed) & jax.random.split(key) — Manages explicit, stateless PRNG keys—always split a key before passing subkeys into independent random draws.

xx → norm → Q/K/V → causal attention → output projection └────────────────────────────── add → residual residual → norm → expansion → GELU → contraction └────────────────────────────────── add → output

3 Run and check

Append this block and run python main.py. Compare its output with your prediction. An assertion failure means a stated contract needs investigation.

main.py · append build step 3
# Step 3 — 3. Run and check: Output shape (4, 8); finite values, zero-projection identity,...
# Create or split explicit PRNG key(s) (`p`) for reproducible randomness.
p = init_block(jax.random.key(7))
# Construct and reshape `x` into the target tensor dimensions.
x = jnp.arange(32, dtype=jnp.float32).reshape(4,8)/10
output = block_forward(p, x)
# Check tensor shape invariant: `output.shape == x.shape`
assert output.shape == x.shape
# Ensure all array elements remain finite: `jnp.all(jnp.isfinite(output))`
assert jnp.all(jnp.isfinite(output))
# Transform every leaf of the parameter PyTree (`zero`).
zero = jax.tree.map(jnp.zeros_like, p)
assert jnp.allclose(block_forward(zero, x), x)
assert jnp.allclose(jax.jit(block_forward)(p,x), output, atol=1e-5)
print("Block shape:", output.shape)

Code building blocks in this step

  • jnp.arange(n, dtype=...) — Creates a 1-D JAX array of evenly spaced values [0, 1, ..., n-1] on the target device.
  • array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
  • jnp.isfinite(x) — Returns a boolean mask verifying that no element is NaN or Inf.

Output shape (4,8)(4, 8); finite values, zero-projection identity, compiled/eager agreement.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build
# Step 1 — 1. Prepare the experiment: These imports provide JAX array transformations and NumPy...
import jax
import jax.numpy as jnp
import numpy as np

# Step 2 — 2. Implement the mechanism: x → norm → Q/K/V → causal attention → output projection...
def layer_norm(x):
    # Reduce along axis=-1 to compute `mean`.
    mean = jnp.mean(x, axis=-1, keepdims=True)
    # Reduce along axis=-1 to compute `variance`.
    variance = jnp.mean((x-mean)**2, axis=-1, keepdims=True)
    return (x-mean) / jnp.sqrt(variance + 1e-5)

def init_block(key, width=8):
    # Create or split explicit PRNG key(s) (`keys`) for reproducible randomness.
    keys = jax.random.split(key, 6)
    shapes = [(width,width)]*4 + [(width,2*width),(2*width,width)]
    return {name: jax.random.normal(k,s)*0.1 for name,k,s in zip(
        ["q","k","v","o","up","down"], keys, shapes)}

def block_forward(p, x):
    h = layer_norm(x)
    q, k, v = h@p["q"], h@p["k"], h@p["v"]
    scores = q@k.T / jnp.sqrt(x.shape[-1])
    # Create evenly spaced index values in `allowed`.
    allowed = jnp.arange(x.shape[0])[:,None] >= jnp.arange(x.shape[0])[None,:]
    # Apply nonlinear activation or probability normalization to compute `weights`.
    weights = jax.nn.softmax(jnp.where(allowed, scores, -jnp.inf), axis=-1)
    residual = x + (weights@v)@p["o"]
    return residual + jax.nn.gelu(layer_norm(residual)@p["up"])@p["down"]

# Step 3 — 3. Run and check: Output shape (4, 8); finite values, zero-projection identity,...
# Create or split explicit PRNG key(s) (`p`) for reproducible randomness.
p = init_block(jax.random.key(7))
# Construct and reshape `x` into the target tensor dimensions.
x = jnp.arange(32, dtype=jnp.float32).reshape(4,8)/10
output = block_forward(p, x)
# Check tensor shape invariant: `output.shape == x.shape`
assert output.shape == x.shape
# Ensure all array elements remain finite: `jnp.all(jnp.isfinite(output))`
assert jnp.all(jnp.isfinite(output))
# Transform every leaf of the parameter PyTree (`zero`).
zero = jax.tree.map(jnp.zeros_like, p)
assert jnp.allclose(block_forward(zero, x), x)
assert jnp.allclose(jax.jit(block_forward)(p,x), output, atol=1e-5)
print("Block shape:", output.shape)

Run the example

python3 phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py
About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example
main.py · complete Python example
# Step 1 — 1. Prepare the experiment: These imports provide JAX array transformations and NumPy...
import jax
import jax.numpy as jnp
import numpy as np

# Step 2 — 2. Implement the mechanism: x → norm → Q/K/V → causal attention → output projection...
def layer_norm(x):
    # Reduce along axis=-1 to compute `mean`.
    mean = jnp.mean(x, axis=-1, keepdims=True)
    # Reduce along axis=-1 to compute `variance`.
    variance = jnp.mean((x-mean)**2, axis=-1, keepdims=True)
    return (x-mean) / jnp.sqrt(variance + 1e-5)

def init_block(key, width=8):
    # Create or split explicit PRNG key(s) (`keys`) for reproducible randomness.
    keys = jax.random.split(key, 6)
    shapes = [(width,width)]*4 + [(width,2*width),(2*width,width)]
    return {name: jax.random.normal(k,s)*0.1 for name,k,s in zip(
        ["q","k","v","o","up","down"], keys, shapes)}

def block_forward(p, x):
    h = layer_norm(x)
    q, k, v = h@p["q"], h@p["k"], h@p["v"]
    scores = q@k.T / jnp.sqrt(x.shape[-1])
    # Create evenly spaced index values in `allowed`.
    allowed = jnp.arange(x.shape[0])[:,None] >= jnp.arange(x.shape[0])[None,:]
    # Apply nonlinear activation or probability normalization to compute `weights`.
    weights = jax.nn.softmax(jnp.where(allowed, scores, -jnp.inf), axis=-1)
    residual = x + (weights@v)@p["o"]
    return residual + jax.nn.gelu(layer_norm(residual)@p["up"])@p["down"]

# Step 3 — 3. Run and check: Output shape (4, 8); finite values, zero-projection identity,...
# Create or split explicit PRNG key(s) (`p`) for reproducible randomness.
p = init_block(jax.random.key(7))
# Construct and reshape `x` into the target tensor dimensions.
x = jnp.arange(32, dtype=jnp.float32).reshape(4,8)/10
output = block_forward(p, x)
# Check tensor shape invariant: `output.shape == x.shape`
assert output.shape == x.shape
# Ensure all array elements remain finite: `jnp.all(jnp.isfinite(output))`
assert jnp.all(jnp.isfinite(output))
# Transform every leaf of the parameter PyTree (`zero`).
zero = jax.tree.map(jnp.zeros_like, p)
assert jnp.allclose(block_forward(zero, x), x)
assert jnp.allclose(jax.jit(block_forward)(p,x), output, atol=1e-5)
print("Block shape:", output.shape)

Expected result: Output shape (4,8)(4, 8); finite values, zero-projection identity, compiled/eager agreement.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:47:18.530225+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

Block shape: (4, 8)
Block shape: (4, 8)
bfloat16 max absolute error: 0.013455629348754883
Float32 reduction policy max absolute error: 0.013455629348754883
PASS: transformers-03

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Stage 3 of 5 · Controlled ExperimentsPredict before running & test boundary conditions

Verify the residual identity and causal boundary

Verify the residual identity and causal boundary · Python experiment
# Experiment — Verify the residual identity and causal boundary: Masking and per-token operations preserve causal independence...
changed_x = x.at[-1].add(100.)
changed_output = block_forward(p, changed_x)
assert jnp.allclose(changed_output[:3], output[:3], atol=1e-5)
assert not jnp.allclose(changed_output[-1], output[-1])

Expected: First three rows agree; the last row changes.

Masking and per-token operations preserve causal independence through both sublayers.

Compare two precision policies

Compare two precision policies · Python experiment
# Experiment — Compare two precision policies: Storage/activation casts and sensitive-reduction precision are...
# Transform every leaf of the parameter PyTree (`low_p`).
low_p = jax.tree.map(lambda z:z.astype(jnp.bfloat16), p)
low_x = x.astype(jnp.bfloat16)
low_out = block_forward(low_p, low_x).astype(jnp.float32)
# Aggregate array values to compute `error`.
error = jnp.max(jnp.abs(low_out-output))
print("bfloat16 max absolute error:", float(error))
# Confirm that all computed values remain finite (no NaN or Inf).
assert jnp.all(jnp.isfinite(low_out))
assert error < 0.1

def mixed_block(p, x):
    h = layer_norm(x.astype(jnp.float32)).astype(jnp.bfloat16)
    q,k,v = h@p["q"], h@p["k"], h@p["v"]
    scores = (q@k.T).astype(jnp.float32)/jnp.sqrt(x.shape[-1])
    # Create evenly spaced index values in `allowed`.
    allowed = jnp.arange(x.shape[0])[:,None] >= jnp.arange(x.shape[0])[None,:]
    weights = jax.nn.softmax(jnp.where(allowed,scores,-jnp.inf),axis=-1).astype(jnp.bfloat16)
    residual = x+(weights@v)@p["o"]
    normalized = layer_norm(residual.astype(jnp.float32)).astype(jnp.bfloat16)
    activation = jax.nn.gelu((normalized@p["up"]).astype(jnp.float32)).astype(jnp.bfloat16)
    return residual + activation@p["down"]

mixed_out = mixed_block(low_p,low_x).astype(jnp.float32)
# Aggregate array values to compute `mixed_error`.
mixed_error = jnp.max(jnp.abs(mixed_out-output))
print("Float32 reduction policy max absolute error:",float(mixed_error))
# Ensure all array elements remain finite: `jnp.all(jnp.isfinite(mixed_out))`
assert jnp.all(jnp.isfinite(mixed_out))
assert mixed_error < 0.1

Expected: Both outputs are finite, with max absolute errors below 0.10.1 for this fixture. Record both actual errors; their relative order is not guaranteed.

Storage/activation casts and sensitive-reduction precision are separate choices. One fixture is distinct from an accuracy or throughput guarantee.

Stage 4 of 5 · Hands-On Exercise & CheckpointModify the code yourself, solve transfer tasks & test your understanding
Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py (or inside the Make it yours cell of exercises/transformers-03.ipynb) and run python3 phases/07-transformers/03-a-transformer-block-and-mixed-precision/code/main.py to verify your prediction before opening the reference solution.

Differentiate the mean squared block output with respect to every projection. Check each gradient leaf’s shape and finite values, then perturb one parameter coordinate and compare a central difference.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
  • jnp.isfinite(x) — Returns a boolean mask verifying that no element is NaN or Inf.
  • jax.grad(loss_fn)(params, ...) — Transforms a scalar-output function into a function returning the gradient PyTree with the same structure as params.

2. Step-by-step implementation plan:

  1. Differentiate the objective to obtain g via automatic differentiation.
  2. Check tensor shape invariant: all(a.shape==b.shape and jnp.all(jnp.isfinite(b)) for a,b in zip(...
Starter code scaffold · copy into your file & fill in the TODOs 
# Exercise solution: Differentiate the mean squared block output with respect to every...
def objective(params):
    return ...  # TODO: return computed result
# Differentiate the objective to obtain `g` via automatic differentiation.
g = jax.grad(...)  # TODO: compute g
# Check tensor shape invariant: `all(a.shape==b.shape and jnp.all(jnp.isfinite(b)) for a,b in zip(...`
assert all(a.shape  # TODO: complete assertion check
def shifted_objective(delta):
    altered = ...  # TODO: compute altered
    return ...  # TODO: return computed result
fd = ...  # TODO: compute fd
assert jnp.allclose(fd,g["down"][0,0],atol=1e-4,rtol=2e-2)  # TODO: complete assertion check
Reveal the reference solution
Reference solution
# Exercise solution: Differentiate the mean squared block output with respect to every...
def objective(params):
    return jnp.mean(block_forward(params,x)**2)
# Differentiate the objective to obtain `g` via automatic differentiation.
g = jax.grad(objective)(p)
# Check tensor shape invariant: `all(a.shape==b.shape and jnp.all(jnp.isfinite(b)) for a,b in zip(...`
assert all(a.shape==b.shape and jnp.all(jnp.isfinite(b)) for a,b in zip(jax.tree.leaves(p),jax.tree.leaves(g)))
def shifted_objective(delta):
    altered = {**p, "down": p["down"].at[0,0].add(delta)}
    return objective(altered)
fd = (shifted_objective(1e-2)-shifted_objective(-1e-2))/(2e-2)
assert jnp.allclose(fd,g["down"][0,0],atol=1e-4,rtol=2e-2)

Check normalization without autodiff

Transfer / diagnosis

Compute the first normalized row using NumPy and compare. Also check a constant row.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
  • x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
  • jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

2. Step-by-step implementation plan:

  1. Aggregate array values to compute reference.
Starter code scaffold · copy into your file & fill in the TODOs 
# Check normalization without autodiff (Transfer / diagnosis): The independent reduction and constant-row case check the...
row = np.asarray(...)  # TODO: compute row
# Aggregate array values to compute `reference`.
reference = ...  # TODO: compute reference
np.testing.assert_allclose(layer_norm(x)[0],reference,atol = ...  # TODO: compute np.testing.assert_allclose(layer_norm(x)[0],reference,atol
assert jnp.allclose(layer_norm(jnp.ones((2,8))),0.)  # TODO: complete assertion check
Show a hint

Use population variance over features and epsilon 10−510^{-5}.

Reveal reference solution and reasoning
Check normalization without autodiff · Python solution
# Check normalization without autodiff (Transfer / diagnosis): The independent reduction and constant-row case check the...
row = np.asarray(x[0])
# Aggregate array values to compute `reference`.
reference = (row-row.mean())/np.sqrt(np.mean((row-row.mean())**2)+1e-5)
np.testing.assert_allclose(layer_norm(x)[0],reference,atol=1e-6)
assert jnp.allclose(layer_norm(jnp.ones((2,8))),0.)

The independent reduction and constant-row case check the normalization mechanism.

Repair a residual width mismatch

Transfer / diagnosis

Remove the contraction from the MLP and reproduce the residual addition failure. Restore the contraction and verify shape.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

  • mismatch(...) — Call mismatch with your updated parameters or inputs from this lesson's workspace.
  • nn.gelu(...) — Call nn.gelu with your updated parameters or inputs from this lesson's workspace.
  • assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

  1. Check tensor shape invariant: expanded.shape == (4,16)
  2. Run the boundary check and catch the expected exception:
  3. Check tensor shape invariant: (x + expanded@p["down"]).shape == x.shape
Starter code scaffold · copy into your file & fill in the TODOs 
# Repair a residual width mismatch (Transfer / diagnosis): Residual paths require matching widths; the down projection...
expanded = jax.nn.gelu(...)  # TODO: compute expanded
# Check tensor shape invariant: `expanded.shape == (4,16)`
assert expanded.shape  # TODO: complete assertion check
# Run the boundary check and catch the expected exception:
try:
    _ = ...  # TODO: compute _
except TypeError:
    pass
else:
    raise AssertionError("expected width mismatch")
# Check tensor shape invariant: `(x + expanded@p["down"]).shape == x.shape`
assert (x + expanded@p["down"]).shape  # TODO: complete assertion check
Show a hint

The expanded tensor has twice the model width.

Reveal reference solution and reasoning
Repair a residual width mismatch · Python solution
# Repair a residual width mismatch (Transfer / diagnosis): Residual paths require matching widths; the down projection...
expanded = jax.nn.gelu(layer_norm(x)@p["up"])
# Check tensor shape invariant: `expanded.shape == (4,16)`
assert expanded.shape == (4,16)
# Run the boundary check and catch the expected exception:
try:
    _ = x + expanded
except TypeError:
    pass
else:
    raise AssertionError("expected width mismatch")
# Check tensor shape invariant: `(x + expanded@p["down"]).shape == x.shape`
assert (x + expanded@p["down"]).shape == x.shape

Residual paths require matching widths; the down projection restores the model-width contract.

Check your understanding

If all projection matrices are zero in this block, what should the output be?

Choose one answer

Diagnose the result

A residual width error means the expanded MLP tensor was added before its contraction. Check (T,2D)(T,2D) versus (T,D)(T,D), restore the down projection and verify the zero-projection identity. For nonfinite or inaccurate lower-precision results, inspect feature variance, epsilon and reduction dtype. Report both measured precision errors rather than assuming a policy must always improve them.

Stage 5 of 5 · Keep Your EvidenceRecord your verification receipt & carry takeaways forward

Carry forward

  • Masking and per-token operations preserve causal independence through both sublayers.
  • The error bound is a checked example, not an accuracy guarantee for arbitrary models.

Keep your evidence

Download evidence template

Keep the architecture diagram, zero-projection identity, causal perturbation, independent normalization, finite-difference coordinate check, precision discrepancy and residual-width repair.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Primary references