Core Foundation · Local CPULaptop CPULesson 2 of 4 in Phase 03

Phase 03 / State, randomness & control flow

Pytrees and structured parameters

Revision v0.7.0 · 2026-10-20

Lesson · about 60 minutes · CPU practice

Represent parameters and state in a nested structure and transform its leaves.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Random keys without surprises
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson 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/03-state/02-pytrees-and-structured-parameters/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/03-state/02-pytrees-and-structured-parameters/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.

Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.

For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter lab

Select the downloaded state-02.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.

Check: Before/after: $4.0$ , approximately $0.64$ .

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

Phase 03.02 Stage 1 of 5 · Understand

1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence

What you will be able to do

Identify containers, leaves and structure
Derive the matching gradient tree
Apply corresponding leaves, not unrelated arrays
Keep metadata separate from differentiated state

The problem

Our model has a weight vector and a bias. How do we keep each parameter paired with its gradient? We’ll store them in a small nested structure called a pytree, calculate one update by hand, and check that JAX follows the same structure. The tree is just a way to keep related values organized.

The idea

A pytree is a nested structure whose leaves hold values such as parameters. JAX can transform the leaves while preserving the structure. This makes a model's parameter groups explicit and lets us compare the gradient tree with the parameters it differentiates.

Structure travels with the leaves

Parameter tree{weight: array, bias: array}
Names and containers describe the structure.
Tree operationmap(update, parameters)
Apply a matching operation to each array leaf.
New parameter treesame structure, new leaves
The optimizer and gradients must agree on the tree.

A pytree is a nested structure with array leaves, not one large flattened tensor. Shape and dtype still matter separately at each leaf.

Conceptual illustration. Verify the computation in the lesson’s Python experiments.

Match leaves by meaning as well as shape

Suppose parameters contain a dense-layer weight matrix and a bias vector. Each gradient leaf describes sensitivity to its corresponding parameter leaf. Match paths such as layer/weight and layer/bias before comparing shapes; two equal-shaped leaves can still be swapped.

Optimizer state is related but need not have exactly the parameter tree's structure. It may contain moment trees plus a scalar step counter or other metadata. A useful diagram aligns matching parameter/gradient paths and then shows the additional optimizer state separately.

When a tree operation fails, inspect structure, leaf paths, shapes and dtypes in that order. Flattening everything immediately may hide the semantic distinction that would explain the error.

Identify containers, leaves and structure

A pytree separates a nested container structure from its leaves. A dictionary with weight and bias keys is a tree; each array is a leaf. A leaf can itself have many numerical elements. Two leaves therefore do not mean two scalar parameters.

JAX transformations understand common dictionaries, lists and tuples. tree.flatten returns an ordered leaf list plus a tree definition that can reconstruct the original containers. Treat that definition as part of the contract, rather than manually guessing where each flattened value belongs.

params
├── weight: float32[2]
└── bias: float32[]
2 leaves, 3 numerical parameter values

Derive the matching gradient tree

Let’s check one example. Take $x=[2,1]$ , weights $w=[1,-1]$ , and bias $b=0$ . The prediction is $1$ , while the target is $3$ , so the residual is $-2$ . For squared loss, the derivative with respect to the prediction is $-4$ . Multiplying by the input gives the weight gradient $-4[2,1]=[-8,-4]$ ; the bias gradient is $-4$ .

With learning rate $0.05$ , the new weights are $[1.4,-0.8]$ , and the bias becomes $0.2$ . The new prediction is $2.2$ , the residual is $-0.8$ , and the loss is $0.64$ . We can compare each of these values with the gradient tree and update. This tells us more than checking only whether the loss decreased.

\begin{aligned}\hat y&=2(1)+1(-1)+0=1\\L&=(1-3)^2=4\\\nabla_w L&=(-8,-4),\qquad \frac{\partial L}{\partial b}=-4\end{aligned}

Apply corresponding leaves, not unrelated arrays

tree.map can apply $p-\eta g$ to corresponding leaves in the parameter and gradient trees. The containers must agree. A missing key is a structural error; an incompatible leaf shape is a numerical contract error. A matching tree structure alone does not prove that each leaf has the correct shape.

Dictionary leaves are traversed in a deterministic key order. Do not zip unrelated dictionaries merely because they have equal lengths. Use the tree definition to preserve identity and add shape checks at model boundaries.

Keep metadata separate from differentiated state

Trainable leaves should be appropriate floating-point arrays. A Python string such as a model name is not a floating-point parameter. Keep such configuration separate, or register a carefully designed custom structure when the model requires it.

Optimizer state is another tree with its own meaning. Its structure need not be identical to parameters. Model parameters, gradients and optimizer statistics are connected, but they are not interchangeable objects. This separation will matter in the Optax lesson and in recovery.

A gradient tree matches parameter leaves

Predict before looking: Which weight changes more, and why?

Read the figure

Each pair of bars compares one parameter before and after a gradient step. Read the categories as the two weight coordinates followed by the bias. The weights move from $(1,-1)$ to $(1.4,-0.8)$ , and the bias moves from $0$ to $0.2$ .

The second weight’s bar remains below zero, but becomes less negative. All three parameters increased; an upward update does not require the parameter itself to be positive.

Connect it to the computation

For the input $(2,1)$ , the initial prediction is $1$ , below the target $3$ . The parameter gradients are $(-8,-4,-4)$ . Subtracting them with learning rate $0.05$ adds $(0.4,0.2,0.2)$ , exactly the changes visible between the paired bars.

The updated prediction is $2.2$ , so squared error falls from $4$ to $0.64$ . The tree structure lets each gradient reach its corresponding parameter leaf. The chart’s vertical axis is parameter value, not loss; you need the prediction calculation to establish improvement.

Open full-size figure Download data and provenance

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

Where to put & run this code: Work in phases/03-state/02-pytrees-and-structured-parameters/code/main.py (or downloaded exercises/state-02.py / exercises/state-02.ipynb) and execute from the repository root with python3 phases/03-state/02-pytrees-and-structured-parameters/code/main.py.

1 Prepare the inputs

Create main.py in your lesson workspace. Add this first block; use the environment from setup.

main.py · append build step 1

# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

These explicit inputs define the case that the later checks will verify.

2 Build the computation

Append this block below the inputs in the same file.

main.py · append build step 2

# Step 2 — Build the computation: loss reads both parameter leaves.
params = {"weight": jnp.array([1., -1.]), "bias": jnp.array(0.)}
x = jnp.array([2., 1.])
def loss(p):
    prediction = jnp.dot(p["weight"], x) + p["bias"]
    return (prediction - 3.) ** 2

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.

loss reads both parameter leaves. grad returns a matching gradient tree; tree.map pairs each parameter with its corresponding derivative for the update.

3 Run and check the result

Append the checks, save main.py, and run python main.py from this folder using your course environment.

main.py · append build step 3

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
# Differentiate the objective to obtain `(value, grads)` via automatic differentiation.
value, grads = jax.value_and_grad(loss)(params)
# Transform every leaf of the parameter PyTree (`updated`).
updated = jax.tree.map(lambda p, g: p - 0.05 * g, params, grads)
print("Before/after:", float(value), float(loss(updated)))
assert grads.keys() == params.keys()
assert loss(updated) < value

Code building blocks in this step

jax.value_and_grad(loss_fn)(params, ...) — Evaluates both the scalar loss and its gradient PyTree (loss_val, grads) in a single forward+backward pass.
jax.tree.map(lambda p, g: ..., params, grads) — Applies a function leaf-by-leaf across matching PyTrees (such as updating every parameter tensor with its gradient).

Compare the output to the expected result below before making the exercise change.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp

# Step 2 — Build the computation: loss reads both parameter leaves.
params = {"weight": jnp.array([1., -1.]), "bias": jnp.array(0.)}
x = jnp.array([2., 1.])
def loss(p):
    prediction = jnp.dot(p["weight"], x) + p["bias"]
    return (prediction - 3.) ** 2

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
# Differentiate the objective to obtain `(value, grads)` via automatic differentiation.
value, grads = jax.value_and_grad(loss)(params)
# Transform every leaf of the parameter PyTree (`updated`).
updated = jax.tree.map(lambda p, g: p - 0.05 * g, params, grads)
print("Before/after:", float(value), float(loss(updated)))
assert grads.keys() == params.keys()
assert loss(updated) < value

Run the example

python3 phases/03-state/02-pytrees-and-structured-parameters/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 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp
# Step 2 — Build the computation: loss reads both parameter leaves.
params = {"weight": jnp.array([1., -1.]), "bias": jnp.array(0.)}
x = jnp.array([2., 1.])
def loss(p):
    prediction = jnp.dot(p["weight"], x) + p["bias"]
    return (prediction - 3.) ** 2
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
# Differentiate the objective to obtain `(value, grads)` via automatic differentiation.
value, grads = jax.value_and_grad(loss)(params)
# Transform every leaf of the parameter PyTree (`updated`).
updated = jax.tree.map(lambda p, g: p - 0.05 * g, params, grads)
print("Before/after:", float(value), float(loss(updated)))
assert grads.keys() == params.keys()
assert loss(updated) < value

Expected result: Before/after: $4.0$ , approximately $0.64$ .

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:45:28.661784+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

Before/after: 4.0 0.6399999260902405
Before/after: 4.0 0.6399999260902405
Expected tree mismatch
PASS: state-02

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

Verify every gradient leaf

Verify every gradient leaf · Python experiment

# Experiment — Verify every gradient leaf: This check verifies parameter identity and the actual gradient...
assert jnp.allclose(grads["weight"],jnp.array([-8.,-4.]))
assert jnp.allclose(grads["bias"],-4.)
assert jnp.allclose(updated["weight"],jnp.array([1.4,-0.8]))
assert jnp.allclose(updated["bias"],0.2)
assert jnp.allclose(loss(updated),0.64,atol=1e-6)

Expected: All leaves agree with the hand-derived update.

This check verifies parameter identity and the actual gradient values, not just a matching container.

Flatten and reconstruct without losing identity

Flatten and reconstruct without losing identity · Python experiment

# Experiment — Flatten and reconstruct without losing identity: The tree definition describes containers; leaves supply...
leaves,definition=jax.tree.flatten(params)
rebuilt=jax.tree.unflatten(definition,leaves)
assert jax.tree.structure(rebuilt)==jax.tree.structure(params)
assert jnp.array_equal(rebuilt["weight"],params["weight"])
assert jnp.array_equal(rebuilt["bias"],params["bias"])

Expected: The original keys and corresponding values are reconstructed.

The tree definition describes containers; leaves supply numerical values. Both are needed for reconstruction.

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/03-state/02-pytrees-and-structured-parameters/code/main.py (or inside the Make it yours cell of exercises/state-02.ipynb) and run python3 phases/03-state/02-pytrees-and-structured-parameters/code/main.py to verify your prediction before opening the reference solution.

Compute the squared norm of all gradient leaves. Verify it by adding the weight and bias contributions directly.

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.

2. Step-by-step implementation plan:

Transform every leaf of the parameter PyTree (norm_squared).
Aggregate array values to compute manual.

Starter code scaffold · copy into your file & fill in the TODOs

# Exercise solution: Compute the squared norm of all gradient leaves.
# Transform every leaf of the parameter PyTree (`norm_squared`).
norm_squared = sum(...)  # TODO: compute norm_squared
# Aggregate array values to compute `manual`.
manual = jnp.sum(...)  # TODO: compute manual
assert jnp.allclose(norm_squared, manual)  # TODO: complete assertion check
assert jnp.allclose(norm_squared, 96.)  # TODO: complete assertion check

Reveal the reference solution

Reference solution

# Exercise solution: Compute the squared norm of all gradient leaves.
# Transform every leaf of the parameter PyTree (`norm_squared`).
norm_squared = sum(jnp.sum(g ** 2) for g in jax.tree.leaves(grads))
# Aggregate array values to compute `manual`.
manual = jnp.sum(grads["weight"] ** 2) + grads["bias"] ** 2
assert jnp.allclose(norm_squared, manual)
assert jnp.allclose(norm_squared, 96.)

Extend to a nested model

Practice

Place weight and bias under a layer dictionary and add a scalar scale parameter. Verify the derivative of scale for $s(wx+b)$ .

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.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
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:

Differentiate the objective to obtain g via automatic differentiation.

Starter code scaffold · copy into your file & fill in the TODOs

# Extend to a nested model (Practice): Nested parameter identity survives differentiation.
nested = ...  # TODO: compute nested
def nested_loss(p):
    pred = ...  # TODO: compute pred
    return ...  # TODO: return computed result
# Differentiate the objective to obtain `g` via automatic differentiation.
g = jax.grad(...)  # TODO: compute g
assert jax.tree.structure(g)  # TODO: complete assertion check
assert jnp.allclose(g["scale"],-4.)  # TODO: complete assertion check

Show a hint

At the starting point prediction is $1$ and residual is $-2$ , so the scale derivative is $-4$ .

Reveal reference solution and reasoning

Extend to a nested model · Python solution

# Extend to a nested model (Practice): Nested parameter identity survives differentiation.
nested={"layer":params,"scale":jnp.array(1.)}
def nested_loss(p):
    pred=p["scale"]*(jnp.dot(p["layer"]["weight"],x)+p["layer"]["bias"])
    return (pred-3.)**2
# Differentiate the objective to obtain `g` via automatic differentiation.
g=jax.grad(nested_loss)(nested)
assert jax.tree.structure(g)==jax.tree.structure(nested)
assert jnp.allclose(g["scale"],-4.)

Nested parameter identity survives differentiation. The independent scale calculation checks the new path.

Repair a mismatched update tree

Challenge

Remove bias from the gradient dictionary, reproduce the structural error, then restore it and verify the hand update.

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.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.tree.map(lambda p, g: ..., params, grads) — Applies a function leaf-by-leaf across matching PyTrees (such as updating every parameter tensor with its gradient).

2. Step-by-step implementation plan:

Run the boundary check and catch the expected exception:
Transform every leaf of the parameter PyTree (fixed).

Starter code scaffold · copy into your file & fill in the TODOs

# Repair a mismatched update tree (Challenge): A missing bias gradient is a tree mismatch.
# Run the boundary check and catch the expected exception:
try:
    jax.tree.map(lambda p,g:p-0.05*g,params,{"weight":grads["weight"]})
except ValueError:
    print("Expected tree mismatch")
else:
    raise AssertionError("Expected a missing-key failure")
# Transform every leaf of the parameter PyTree (`fixed`).
fixed = jax.tree.map(...)  # TODO: compute fixed
assert jnp.allclose(loss(fixed),0.64,atol=1e-6)  # TODO: complete assertion check

Show a hint

Matching leaf counts is not enough; key paths matter.

Reveal reference solution and reasoning

Repair a mismatched update tree · Python solution

# Repair a mismatched update tree (Challenge): A missing bias gradient is a tree mismatch.
# Run the boundary check and catch the expected exception:
try:
    jax.tree.map(lambda p,g:p-0.05*g,params,{"weight":grads["weight"]})
except ValueError:
    print("Expected tree mismatch")
else:
    raise AssertionError("Expected a missing-key failure")
# Transform every leaf of the parameter PyTree (`fixed`).
fixed=jax.tree.map(lambda p,g:p-0.05*g,params,grads)
assert jnp.allclose(loss(fixed),0.64,atol=1e-6)

A missing bias gradient is a tree mismatch. Broadcasting weights or changing a learning rate cannot supply the missing parameter identity.

Diagnose the result

A missing bias gradient is a tree mismatch. Broadcasting weights or changing a learning rate cannot supply the missing parameter identity.

Carry forward

This check verifies parameter identity and the actual gradient values, not just a matching container.
The tree definition describes containers; leaves supply numerical values. Both are needed for reconstruction.

Keep your evidence

Download evidence template

Keep the parameter/gradient tree diagram, hand-derived leaves and update, flatten/unflatten round trip, nested-scale derivative and missing-key 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

JAX pytrees ↗

Pytrees and structured parameters

What you will be able to do

The problem

The idea

Match leaves by meaning as well as shape

Identify containers, leaves and structure

Derive the matching gradient tree

Apply corresponding leaves, not unrelated arrays

Keep metadata separate from differentiated state

A gradient tree matches parameter leaves

Read the figure

Connect it to the computation

1 Prepare the inputs

2 Build the computation

3 Run and check the result

Run the example

Recorded reference execution

Verify every gradient leaf

Flatten and reconstruct without losing identity

Make it yours

Extend to a nested model

Repair a mismatched update tree

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references