Core Foundation · Local CPULaptop CPULesson 3 of 4 in Phase 01

Phase 01 / Arrays & pure functions

Pure functions and explicit inputs

Revision v0.7.0 · 2026-10-20

Lesson · about 45 minutes · CPU practice

Design an explicit parameter/data interface and return new state without mutating old containers.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Shapes, broadcasting, and dtypes
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/01-arrays/03-pure-functions-and-explicit-inputs/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/01-arrays/03-pure-functions-and-explicit-inputs/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 arrays-03.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: [ $1$ . $3$ . $5$ .]

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

Phase 01.03 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 hidden inputs and side effects in a numerical function.
Pass model parameters explicitly and return new state without mutating input containers.
Test replay and independence between two parameter configurations.
Explain why host orchestration and transformed numerical computation need different boundaries.

The problem

Have you ever rerun a notebook cell and got a different result because another cell changed a variable? Let’s make that dependency visible. We’ll pass a model’s weights into its prediction function, return new values explicitly, and check that the same inputs reproduce the same calculation.

The idea

A pure function makes its changing inputs visible and returns its result without secretly changing shared state. That gives transformations such as differentiation and compilation a clear computation to work with. When a prediction changes, we should be able to identify the input that changed.

Make a repeated call explainable

Imagine a prediction function that receives features and weights but reads a global bias. You compile it, change that Python global, and call the compiled function again. The source file now displays the new bias, but the already-traced computation may still embody the old value. The hidden dependency makes the result difficult to reason about.

Pass the bias as an explicit argument instead. Then two calls with different bias arrays are visibly two computations with different data. For a linear prediction, increasing the bias by $1$ must increase every output by $1$ , while changing the weight vector can affect rows differently.

Purity does not mean a whole application has no state. It means the numerical function receives the state it needs and returns any new state. Files, logs and orchestration can live outside that numerical boundary.

X'[i] = \begin{cases} v & \text{if } i = k \\ X[i] & \text{otherwise} \end{cases}

Hidden state makes the signature incomplete

A function that takes $x$ but reads a global weight has two effective inputs while advertising only one. Re-running it after changing the global can change its result with no visible argument change. Tests may depend on the order in which prior code ran.

Passing the weight explicitly makes the dependency visible. It also allows two model configurations to coexist without rewriting global variables. That matters for evaluation, comparisons, checkpoint recovery, and eventually multiple workers.

Array immutability does not make a dictionary immutable

JAX arrays are immutable values, but a Python dictionary holding them is still a mutable container. Assigning params["weight"] = ... changes which value the dictionary holds. If another variable refers to the same dictionary, it sees the changed entry.

Constructing {**params, "weight": new_weight} makes a new dictionary with a changed binding and keeps the old container intact. This is a shallow copy: nested mutable objects can still be shared. Our example has only array leaves, so the shallow update is sufficient for the declared structure. More general parameter trees will need deliberate tree-aware transformations.

Treat a step as an input/output contract

A training step will eventually consume parameters, optimizer state, and a batch, then return updated state and metrics. The current predictor is the simpler first piece: parameters and data in, predictions out. The caller owns the history and can save both original and updated values.

This structure makes replay a meaningful test. With the same inputs, predictions should agree; changing the weight should affect predictions in the expected way. Do not make a test depend on a global that may already have been edited elsewhere in a notebook.

predict(parameters, observations) → predictions
future step(state, batch) → new_state, metrics

Choose a clear transformation boundary

A print, file write, or list append is a Python side effect. Inside a transformed function, such effects can happen during tracing rather than every device execution. We will investigate that mechanism in the tracing lesson. For now, separate producing metrics from reporting them.

Return the values you need to inspect, then let the caller print or save them. This design also makes unit-sized numerical checks simpler: you can compare actual outputs with known arrays without parsing console messages or relying on side-effect timing.

A deterministic check has limits

Calling a function twice with fixed inputs checks a useful property, but it does not prove the absence of every hidden dependency. A hidden global might happen not to change. Inspect the implementation and deliberately exercise multiple configurations. Randomness will need an explicit input as well; that is why the next phase introduces keys rather than treating random draws as invisible global state.

Keep the evidence proportional to the claim: these tests establish the behavior of this predictor and update pattern, not the replay of an entire distributed training job.

Explicit parameters make predictions inspectable

Predict before looking: What changes when the bias increases by $2$ ?

Read the figure

The horizontal axis is the input and the vertical axis is the prediction. The solid line uses bias $1$ ; the dashed line uses bias $3$ . Both use the same weight $2$ .

At input $0$ , the predictions are $1$ and $3$ . At input $2$ , they are $5$ and $7$ . The gap stays $2$ everywhere, and the parallel lines rise by $2$ for each unit increase in the input.

Connect it to the computation

In $y=wx+b$ , changing the explicit bias changes the vertical offset, while the unchanged weight preserves the slope. You can read the bias directly where the line meets the vertical axis because the product $wx$ is zero there.

This makes the function’s inputs inspectable: the change in the output follows the change in the passed parameter. As a check, setting the bias to $0$ would move the solid line down by $1$ and make it pass through the origin. It would not flatten the line.

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/01-arrays/03-pure-functions-and-explicit-inputs/code/main.py (or downloaded exercises/arrays-03.py / exercises/arrays-03.ipynb) and execute from the repository root with python3 phases/01-arrays/03-pure-functions-and-explicit-inputs/code/main.py.

1 Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for pure functions and explicit inputs.

main.py · append build step 1

import jax.numpy as jnp

Establishing explicit input shapes and dtypes first makes the downstream transformation contract deterministic.

2 Step 2: Apply the core JAX transformation

Write the core computation and transformation step over the initialized inputs.

main.py · append build step 2

def predict(params, x):
    return params["weight"] * x + params["bias"]
params = {"weight": jnp.array(2.), "bias": jnp.array(1.)}

Code building blocks in this step

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

This stage executes the primary numerical transformation and binds the intermediate outputs.

3 Step 3: Verify shapes and numerical invariants

Check that the resulting arrays satisfy the expected shape, dtype, and numerical tolerances.

main.py · append build step 3

x = jnp.array([0., 1., 2.])
assert jnp.allclose(predict(params, x), jnp.array([1., 3., 5.]))

Code building blocks in this step

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.

These assertions lock in the exact numerical contract before you run the full experiment and variations.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

import jax.numpy as jnp

def predict(params, x):
    return params["weight"] * x + params["bias"]
params = {"weight": jnp.array(2.), "bias": jnp.array(1.)}

x = jnp.array([0., 1., 2.])
assert jnp.allclose(predict(params, x), jnp.array([1., 3., 5.]))

Run the example

python3 phases/01-arrays/03-pure-functions-and-explicit-inputs/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

# Pure functions and explicit inputs: A pure function makes its changing inputs visible and returns...
import jax.numpy as jnp
def predict(params, x):
    return params["weight"] * x + params["bias"]
params = {"weight": jnp.array(2.), "bias": jnp.array(1.)}
x = jnp.array([0., 1., 2.])
print(predict(params, x))
assert jnp.allclose(predict(params, x), jnp.array([1., 3., 5.]))

Expected result: [ $1$ . $3$ . $5$ .]

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

[1. 3. 5.]
aliased container weight: 9.0
evaluation losses: 0.0 6.0
PASS: arrays-03

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

Expose a mutable container alias

Expose a mutable container alias · Python experiment

# Experiment — Expose a mutable container alias: The mutation concerns Python container bindings.
container = {"weight": jnp.array(2.), "bias": jnp.array(1.)}
alias = container
alias["weight"] = jnp.array(9.)
assert float(container["weight"]) == 9.
assert float(params["weight"]) == 2.
print("aliased container weight:", float(container["weight"]))

Expected: The aliased container changes to weight $9$ ; the separate original lesson parameters remain $2$ .

The mutation concerns Python container bindings. It does not contradict JAX array immutability.

Replay two independent model configurations

Replay two independent model configurations · Python experiment

# Experiment — Replay two independent model configurations: Independent configurations reveal the benefit of explicit inputs...
first_params = {"weight": jnp.array(2.), "bias": jnp.array(1.)}
second_params = {"weight": jnp.array(-1.), "bias": jnp.array(4.)}
first_output = predict(first_params, x)
second_output = predict(second_params, x)
assert jnp.allclose(first_output, jnp.array([1., 3., 5.]))
assert jnp.allclose(second_output, jnp.array([4., 3., 2.]))
assert jnp.allclose(predict(first_params, x), first_output)
assert jnp.allclose(predict(second_params, x), second_output)

Expected: The two configurations produce $[1, 3, 5]$ and $[4, 3, 2]$ , and both can be replayed.

Independent configurations reveal the benefit of explicit inputs more clearly than one call on one global parameter set.

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/01-arrays/03-pure-functions-and-explicit-inputs/code/main.py (or inside the Make it yours cell of exercises/arrays-03.ipynb) and run python3 phases/01-arrays/03-pure-functions-and-explicit-inputs/code/main.py to verify your prediction before opening the reference solution.

Create new parameters with weight $3$ while retaining the original bias. Show that the original parameters still produce the original predictions.

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:

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.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

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

# Exercise solution: Create new parameters with weight 3 while retaining the original bias.
updated = ...  # TODO: compute updated
assert jnp.allclose(predict(updated, x), jnp.array([1.,4.,7.]))  # TODO: complete assertion check
assert float(params["weight"])  # TODO: complete assertion check

Reveal the reference solution

Reference solution

# Exercise solution: Create new parameters with weight 3 while retaining the original bias.
updated = {**params, "weight": jnp.array(3.)}
assert jnp.allclose(predict(updated, x), jnp.array([1.,4.,7.]))
assert float(params["weight"]) == 2.0

Update the bias without changing either parameter input

Practice

Implement with_bias(params, new_bias), returning a new dictionary. Verify outputs for new bias $-2$ and replay the original prediction.

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.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

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

# Update the bias without changing either parameter input (Practice): The caller can evaluate both the old and new model without...
def with_bias(parameters, new_bias):
    return ...  # TODO: return computed result
changed_bias = with_bias(...)  # TODO: compute changed_bias
assert changed_bias is not params  # TODO: complete assertion check
assert jnp.allclose(predict(changed_bias, x), jnp.array([-2., 0., 2.]))  # TODO: complete assertion check
assert jnp.allclose(predict(params, x), jnp.array([1., 3., 5.]))  # TODO: complete assertion check

Show a hint

A new dictionary is required; changing a binding in the original dictionary breaks the promised contract.

Reveal reference solution and reasoning

Update the bias without changing either parameter input · Python solution

# Update the bias without changing either parameter input (Practice): The caller can evaluate both the old and new model without...
def with_bias(parameters, new_bias):
    return {**parameters, "bias": jnp.asarray(new_bias)}
changed_bias = with_bias(params, -2.)
assert changed_bias is not params
assert jnp.allclose(predict(changed_bias, x), jnp.array([-2., 0., 2.]))
assert jnp.allclose(predict(params, x), jnp.array([1., 3., 5.]))

The caller can evaluate both the old and new model without restoring a global value or depending on call order.

Return predictions and metrics together

Challenge

Write evaluate(params, inputs, targets) that returns predictions and a scalar MSE. Test two model configurations on targets $[1, 3, 5]$ . Keep printing outside the function.

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.
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:

Check tensor shape invariant: error.shape == ()

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

# Return predictions and metrics together (Challenge): The second model has residuals [3, 0, -3], giving mean...
def evaluate(parameters, inputs, targets):
    predictions = predict(...)  # TODO: compute predictions
    return ...  # TODO: return computed result
labels = jnp.array(...)  # TODO: compute labels
outputs, error = evaluate(...)  # TODO: compute outputs, error
other_outputs, other_error = evaluate(...)  # TODO: compute other_outputs, other_error
# Check tensor shape invariant: `error.shape == ()`
assert error.shape  # TODO: complete assertion check
assert jnp.allclose(error, 0.)  # TODO: complete assertion check
assert jnp.allclose(other_error, 6.)  # TODO: complete assertion check
assert jnp.allclose(outputs, first_output)  # TODO: complete assertion check
print("evaluation losses:", float(error), float(other_error))

Show a hint

Compute data inside the numerical core and report it outside. A tuple return does not itself imply mutation.

Reveal reference solution and reasoning

Return predictions and metrics together · Python solution

# Return predictions and metrics together (Challenge): The second model has residuals [3, 0, -3], giving mean...
def evaluate(parameters, inputs, targets):
    predictions = predict(parameters, inputs)
    return predictions, jnp.mean((predictions - targets) ** 2)
labels = jnp.array([1., 3., 5.])
outputs, error = evaluate(first_params, x, labels)
other_outputs, other_error = evaluate(second_params, x, labels)
# Check tensor shape invariant: `error.shape == ()`
assert error.shape == ()
assert jnp.allclose(error, 0.)
assert jnp.allclose(other_error, 6.)
assert jnp.allclose(outputs, first_output)
print("evaluation losses:", float(error), float(other_error))

The second model has residuals $[3, 0, -3]$ , giving mean squared error $6$ . Known output arithmetic makes the check independent of the implementation.

Diagnose the result

If a replay differs, inspect hidden globals, mutable aliases, and random state before assuming nondeterministic hardware. If an update affects an old model, check whether the caller reused the same dictionary. If logs disappear under a transform, inspect whether they are Python side effects. Preserve numerical results as return values so you can investigate without relying on side-effect behavior.

Carry forward

Explicit inputs make model dependencies visible and configurations independent.
Immutable arrays can still be stored in mutable Python containers.
Return new state deliberately and let the caller own it.
Keep numerical computation separate from host logging, I/O, and orchestration.

Keep your evidence

Download evidence template

Keep both parameter configurations, their known predictions, a container-alias experiment, a non-mutating update, and an evaluation loss checked by hand. Explain which state the caller owns.

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).

Pure functions and explicit inputs

What you will be able to do

The problem

The idea

Make a repeated call explainable

Hidden state makes the signature incomplete

Array immutability does not make a dictionary immutable

Treat a step as an input/output contract

Choose a clear transformation boundary

A deterministic check has limits

Explicit parameters make predictions inspectable

Read the figure

Connect it to the computation

1 Step 1: Set up imports and input tensors

2 Step 2: Apply the core JAX transformation

3 Step 3: Verify shapes and numerical invariants

Run the example

Recorded reference execution

Expose a mutable container alias

Replay two independent model configurations

Make it yours

Update the bias without changing either parameter input

Return predictions and metrics together

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references