Chapter 1 · welcome-01

Set up your learning workspace

Download a starter workspace, select its Python environment, run a saved program, and verify one deliberate change.

What you will be able to do

• Open a terminal in the folder containing your experiment.
• Create a project-specific Python environment and check which interpreter you are using.
• Install the course packages and run a saved Python file.
• Recognize a successful result, save the output, and resume in a new terminal.

The problem

Let’s get your first JAX experiment running. You do not need to know Git or terminal commands yet: we will show where to type each command and how to check the result. You will need a computer, an internet connection for the downloads, and permission to install Python. By the end, you’ll have one experiment and a record of what it produced. A laptop CPU is enough.

The idea

Your first goal is a small, repeatable result: save a Python file, run it with the intended Python environment, and explain its output. A terminal accepts commands; Python executes the program; JAX is a package that Python imports. Keeping those three roles separate makes setup errors much easier to locate.

Follow a file from saving to running

Imagine that your editor shows first_experiment.py, but the terminal says the file does not exist. The program has not reached JAX yet. First save the file, check its full filename, and compare its folder with the terminal's working folder. A file visible in an editor can still be unsaved or stored elsewhere.

Now suppose Python finds the file but cannot import JAX. That is a different boundary: the selected interpreter cannot find the package. Check the interpreter path and its environment before installing anything again. On Windows, the commands below select the environment's Python directly; shell activation is not required.

Once the calculation runs, retain the command and printed result together. Tomorrow, the command tells you how to reproduce the result; a screenshot of the number alone cannot identify the file or environment that produced it.

Where the first result comes from

Predict: The notebook imports JAX, but the terminal script does not. What should you compare first?

Read downward: a command runs from a working folder, selects an interpreter, opens the file and imports its packages before computing. Arrows are dependencies, not measured time. A file-path failure happens before JAX arithmetic; an import failure points to the selected environment.

1. Download and unpack your workspace

Use the “Download the beginner workspace” button above. Open your Downloads folder and unzip jax-start-here.zip. On Windows, right-click it and choose Extract All; on macOS, double-click it. Open the extracted jax-start-here folder. Do not work inside the ZIP preview.

You should see first_experiment.py, requirements-cpu.txt, and START-HERE.md. first_experiment.py is the program you will run. requirements-cpu.txt lists the package versions to install. START-HERE.md is a copy of these instructions. Keep your own files alongside them, outside .venv. You do not need to clone the repository for this lesson.

2. Check that Python is installed

Open a terminal. On macOS, press Command–Space, type Terminal, and press Return. On Windows, open Start, type PowerShell, and open it. On Linux, open your desktop’s Terminal application.

Choose the command for your operating system below. Type or paste one line, then press Enter. Do not type a prompt symbol such as $ or PS>. If you see >>>, you are already inside the Python interpreter; type exit() and press Enter to return to the terminal before using these commands.

The course CPU examples were checked with Python 3.14.3. These commands select Python $3.14$. If you get “command not found” or “not recognized,” Python $3.14$ is not available through that command yet. Open the official Python downloads link in the references, select a Python $3.14$ installer for macOS or Windows, and follow the installer. On Windows, the launcher/install manager supplies py; on macOS, the Python installer supplies python3.14. Close and reopen the terminal after installing, then retry the version check. Linux installation depends on the distribution: use its Python $3.14$ installation instructions and return here once the version check works. Do not continue by substituting an old system Python without checking compatibility.

macOS / Linux — terminal

# Run macos / linux — terminal using the course Python environment
python3.14 --version

Expected:

A line beginning Python 3.14. The last patch number may differ from the tested 3.14.3.

Windows — PowerShell

# Windows — PowerShell
py -3.14 --version

Expected:

A line beginning Python 3.14. If py exists but cannot find 3.14, install that Python version first.

3. Put the terminal in the right folder

The program and requirements file are in the extracted workspace. A command using a filename looks for it relative to the terminal’s current folder. We need the terminal to be in jax-start-here.

macOS: type cd followed by a space, drag the extracted jax-start-here folder from Finder into the terminal, then press Enter. Dragging inserts its actual path, including spaces correctly. Linux: use cd with the actual folder path in quotes; your file manager may offer “Open in Terminal.” Windows: open the extracted folder in File Explorer, click its address bar, type powershell, and press Enter. That opens PowerShell in this folder.

Run the appropriate check below. If requirements-cpu.txt is missing, move into the extracted folder containing that file before continuing. “Repository root” in later lessons means the top folder of the full course checkout; this beginner workspace is enough for the current experiment.

macOS / Linux — check folder and files

# macOS / Linux — check folder and files
pwd
ls

Expected:

The path ends with your extracted workspace folder, and the file list includes first_experiment.py and requirements-cpu.txt.

Windows — check folder and files

# Windows — check folder and files
Get-Location
Get-ChildItem

Expected:

The displayed folder contains first_experiment.py and requirements-cpu.txt.

4. Create the environment and install the packages

Run the lines for your operating system in order, staying in the workspace folder. The venv command creates .venv using the Python version you just selected. It often finishes without printing anything. That is normal; wait until the terminal is ready for another command. If you already have a .venv from another Python version in this folder, use a fresh extracted workspace for this walkthrough.

On macOS/Linux, source loads the activation script into this terminal. You may see (.venv) at the start of the prompt. The important check is sys.executable: it should point inside this workspace’s .venv/bin. On Windows we use .venv\Scripts\python.exe explicitly, so no activation script or execution-policy change is needed.

The pip command is the package installer. -$r$ means “read the package list from this file.” It downloads and installs packages into the selected environment. Let it finish. Download messages are normal; a traceback or an ERROR means the step did not succeed. The final import check below is more useful than guessing from the download log. The pins describe the tested CPU course setup; accelerator installation is a later lesson.

macOS / Linux — run one line at a time

# Run macos / linux — run one line at a time using the course Python environment
python3.14 -m venv .venv
source .venv/bin/activate
python -c "import sys; print(sys.executable)"
python -m pip install -r requirements-cpu.txt
python -c "import jax, numpy, optax; print(jax.__version__, numpy.__version__, optax.__version__)"

Expected:

The interpreter path includes .venv/bin/python. The final line is 0.9.2 2.4.4 0.2.8.

Windows — run one line at a time

# Windows — run one line at a time
py -3.14 -m venv .venv
.\.venv\Scripts\python.exe -c "import sys; print(sys.executable)"
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt
.\.venv\Scripts\python.exe -c "import jax, numpy, optax; print(jax.__version__, numpy.__version__, optax.__version__)"

Expected:

The interpreter path includes .venv\Scripts\python.exe. The final line is 0.9.2 2.4.4 0.2.8. These PowerShell instructions have not been executed on Windows in this project.

5. Run your first saved program

Open first_experiment.py in a text/code editor if you want to see its contents. The same Python program is shown below in “Run the example.” Python code belongs in that file; the commands below belong in the terminal. The downloaded file is already saved, so you can run it without copying any code.

Run the command for your operating system. The program prints the Python and package versions, the backend, the devices, and a sum. The starter explicitly selects CPU before importing JAX, so the expected backend here is cpu. The device description can vary; the final result should be Sum: $6.0$. The program constructs the numbers $0$, $1$, $2$, $3$; adding them gives $6$. This known result checks that we ran a numerical computation, not only installed a package.

If you see a traceback, read the final line and compare it with the troubleshooting section. Do not count a failed command as a completed experiment.

macOS / Linux — activated terminal

# Run macos / linux — activated terminal using the course Python environment
python first_experiment.py

Expected:

Python: 3.14.x JAX: 0.9.2 NumPy: 2.4.4 Backend: cpu Devices: a CPU device description Sum: 6.0

Windows — PowerShell

# Run windows — powershell using the course Python environment
.\.venv\Scripts\python.exe first_experiment.py

Expected:

The same version/backend/result lines. CPU device wording may differ.

6. Make one change and save your result

In the editor, find arange(4, ...) in first_experiment.py and change $4$ to $6$. The array now contains $0$ through $5$. Predict their sum before running. Also change the final assert from == 6.0 to == 15.0, because that assertion is the program’s expected-result check. Save the file with its .py extension; a filename such as first_experiment.py.txt will not match the run command.

Run it again using the Step $5$ command. The final line should be Sum: $15.0$. If it still says $6.0$, check that you saved the file and that the terminal is in the same workspace folder you edited.

After that successful run, the commands below run it once more and save its printed output to my-first-run.txt. Open that text file and check that the versions, CPU information, and sum are present. Add a note in your editor explaining the one change and your prediction. The > symbol saves output to a file and replaces its previous contents; choose another filename when you want to keep multiple runs.

macOS / Linux — save output

# Run macos / linux — save output using the course Python environment
python first_experiment.py > my-first-run.txt

Expected:

my-first-run.txt appears in the workspace and ends with Sum: 15.0 after your saved edit.

Windows — save output

# Run windows — save output using the course Python environment
.\.venv\Scripts\python.exe first_experiment.py > my-first-run.txt

Expected:

my-first-run.txt contains the program output after your saved edit.

7. Resume when you open a new terminal

The .venv folder stays on disk when you close the terminal. On macOS/Linux, activation only applied to the terminal you used; a new terminal needs it again. Move into the workspace folder using Step $3$, then use the commands below. You do not need to reinstall the packages on every session.

On Windows, move into the workspace and use the environment’s full Python path again. To leave an activated macOS/Linux environment, type deactivate. That changes the terminal’s command selection; it does not delete your code or installed packages.

macOS / Linux — new terminal, same folder

# macOS / Linux — new terminal, same folder
source .venv/bin/activate
python first_experiment.py

Expected:

The same saved program runs using the existing environment.

Windows — new PowerShell, same folder

# Run windows — new powershell, same folder using the course Python environment
.\.venv\Scripts\python.exe first_experiment.py

Expected:

The explicit environment interpreter runs the saved program without activation.

When a step fails: match the symptom

“No such file” / “Cannot find path” for requirements-cpu.txt: check Step $3$. The file must be in the current folder. Do not create an empty replacement file.

“No module named jax”: first check which Python is running. Repeat the interpreter-path check in Step $4$, then install using that exact environment interpreter. A different terminal or editor may be selecting a different environment.

“No module named venv” or a message that ensurepip is unavailable: your Python distribution is missing its environment component. On Linux, use the distribution’s matching Python/venv package instructions; package names vary. This needs fixing before the environment can be created.

“No matching distribution found” or a failed jaxlib wheel download: record the Python version, operating system, CPU architecture, and full error. Check the JAX installation support table linked below. Do not remove the version pins or switch to TPU packages as a blind workaround. CPU platform support is not identical across operating systems; this repository’s Windows walkthrough is documented but not execution-validated.

A SyntaxError after pasting a terminal command: check whether you pasted it into a Python file or at the >>> prompt. Terminal commands and Python statements have different places to run.

Sum: $15.0$ followed by AssertionError: you changed the array length but left the old expected sum in the assertion. Update the expectation after deriving the new result; do not delete the check merely to hide the failure.

Workspace array verification across input elements

Predict before looking: Before running, predict the element values and squared values of the 1D JAX array x.

Read the figure

Each index shows the initialized JAX array element alongside its elementwise square on the active device.

Connect it to the computation

Running this check confirms that JAX, NumPy, and Matplotlib execute deterministically in your local environment.

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

Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for set up your learning workspace.

import os
# Choose CPU before importing JAX for this first experiment.
os.environ["JAX_PLATFORMS"] = "cpu"

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

Step 2: Apply the core JAX transformation

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

import platform
import jax

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

Step 3: Verify shapes and numerical invariants

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

import numpy as np
x = jax.numpy.arange(4, dtype=jax.numpy.float32)
assert float(x.sum()) == 6.0

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

Run the example

# Set up your learning workspace: Your first goal is a small, repeatable result: save a Python...
import os
# Choose CPU before importing JAX for this first experiment.
os.environ["JAX_PLATFORMS"] = "cpu"
# Import required JAX, NumPy, and standard-library modules.
import platform
import jax
import numpy as np
print("Python:", platform.python_version())
print("JAX:", jax.__version__)
print("NumPy:", np.__version__)
print("Backend:", jax.default_backend())
print("Devices:", jax.devices())
x = jax.numpy.arange(4, dtype=jax.numpy.float32)
print("Sum:", float(x.sum()))
assert float(x.sum()) == 6.0

Python: 3.14.3
JAX: 0.9.2
NumPy: 2.4.4
Backend: cpu
Devices: [CpuDevice(id=0)]
Sum: 6.0
Changed experiment sum: 15.0
PASS: welcome-01

Your exercise

Edit the starter as shown in Step $6$: change the array length from $4$ to $6$ and the expected sum from $6$ to $15$. Save, rerun, and keep my-first-run.txt plus a note explaining why the result changed.

# Exercise solution: Edit the starter as shown in Step 6: change the array length from 4 to...
x = jax.numpy.arange(6, dtype=jax.numpy.float32)
print("Changed experiment sum:", float(x.sum()))
assert float(x.sum()) == 15.0

Checkpoint

Which evidence lets another learner reproduce your environment?

1. A screenshot of the final number only
2. Python/package versions, device information, and the script
3. The name of your editor

Diagnose the result

Use the troubleshooting section to identify the failing step. Keep the command and final error line, then verify the fix by repeating that step. Successful package installation, successful imports, and a correct numerical result are three separate checks.

Carry forward

• Terminal commands run from a working folder; Python statements live in the .py file.
• The environment interpreter determines which installed packages the program uses.
• Expected values let you check the numerical result after a deliberate edit.
• Keep your code and results outside .venv so the environment can be recreated.

Your evidence

Keep first_experiment.py, my-first-run.txt with package/device output, the environment interpreter path, and a note deriving the change from a sum of 6 to 15. Record a setup error and its repair if you encountered one.

Primary references

• Python downloads: choose Python 3.14 for your operating system
• Python: creating and using virtual environments
• JAX: installation and supported platforms

Chapter 2 · welcome-02

Meet your arrays and devices

Create an array, inspect its dtype and device, and explain the result.

What you will be able to do

• Read the table before using an API
• Derive a reduction and its output shape
• Distinguish dtype from mathematical value
• Inspect placement without assuming an accelerator

The problem

Imagine a small table: two observations, with three measurements in each row. If we ask JAX to add the values, should it add across the rows or down the columns? Let’s work out both answers before running any code. You’ll learn to check an array’s shape, number format and device, so a plausible-looking result does not hide a mistake.

The idea

An array is a collection of values with a shape, a dtype and a device. Shape tells us how to index the values; dtype tells us how they are represented; device tells us where the computation lives. Start by predicting a small reduction before asking JAX to perform it.

Read one array two ways

Take the two rows $[1,2,3]$ and $[4,5,6]$. Summing across each row gives $[6,15]$: one answer per observation. Summing down the rows gives $[5,7,9]$: one answer per feature. The values are the same; the question determines the axis.

In Python, the row sums use axis=1, because the feature axis is reduced. The feature sums use axis=0, because the observation axis is reduced. The axis you name is the one you consume, not the one you keep.

Read the heatmap by locating one row and one column before looking at its color. Color helps compare values, but labels establish which observation and feature you are comparing. A device label is a separate property; moving the same values to another device does not change what a row means.

Read the table before using an API

Start with the six values $0$ through $5$. Reshape them into two rows: $[0, 1, 2]$ and $[3, 4, 5]$. The first axis selects an observation; the second selects a measurement. Shape $(2, 3)$ describes that organization, not the numerical values. Reshape preserves element count and ordering here; it does not create new measurements.

A length-six array and a two-by-three array contain the same values but invite different operations. In your notebook, write an axis label beside each dimension. This becomes important when a later model expects one prediction per row.

flat [0,1,2,3,4,5] → rows [[0,1,2],[3,4,5]]
axis 0: observations (2)
axis 1: measurements (3)

Derive a reduction and its output shape

Summing along axis $1$ removes the measurement axis and leaves one number per observation. The first row sums to $3$, the second to $12$, so the result has shape $(2,)$. Summing along axis $0$ instead combines observations and leaves three measurement totals: $[3, 5, 7]$.

A reduction can also keep the reduced axis with keepdims=True. Column means then have shape $(1, 3)$, which states explicitly that the same three means are subtracted from every row. Derive the numbers and shape before using sum or mean.

(2,3) --sum(axis=1)--> (2,) : [3,12]
(2,3) --sum(axis=0)--> (3,) : [3,5,7]

Distinguish dtype from mathematical value

Dtype describes storage and arithmetic, not whether a value happens to be an integer. The value $3$ can be represented as int32 or float32. This lesson asks for float32 explicitly so later differentiation uses an inexact input. Do not infer precision from the way Python prints a scalar.

Float32 cannot represent every real number, or even every large integer. An exact-looking printed value is not a guarantee of exact arithmetic. Keep dtype in an experiment record and compare with tolerances when arithmetic is approximate.

Inspect placement without assuming an accelerator

jax.devices() lists available devices; x.devices() describes the placement of this array. A CPU result is useful evidence of CPU execution. It is separate from that a TPU runtime is installed or that a model scales across devices.

Creating a NumPy copy brings values into host memory for inspection. Converting a scalar with float also requires a concrete value. Those are useful boundaries for tests and reports; repeated transfers inside a training loop can distort performance. Use device metadata to inspect placement and an explicit synchronization point for timing later.

Rows are groups of observations

Predict before looking: Which row should sum to $12$?

Read the figure

The horizontal axis selects a feature column; the vertical axis selects an observation row. Each cell contains one array value. The color bar maps larger numbers to darker purple: the bottom-right cell is darkest because it contains $5$.

Read across the top row to find $(0,1,2)$, then across the bottom row to find $(3,4,5)$. The picture therefore represents an array with shape $(2,3)$: two observations, each with three features.

Connect it to the computation

Now connect the picture to a reduction. Summing the feature columns within each row gives $0+1+2=3$ and $3+4+5=12$. That is why summing along axis $1$ returns two values. Summing down rows instead would return $(3,5,7)$, one value per feature.

The row labels are positions, not times or class labels. The darker lower row tells us its stored values are larger; it does not imply that the observation is more important.

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

Prepare the inputs

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

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

Build the computation

Append this block below the inputs in the same file.

# Step 2 — Build the computation: x has two observations and three features.
# Construct and reshape `x` into the target tensor dimensions.
x = jnp.arange(6, dtype=jnp.float32).reshape(2, 3)

$x$ has two observations and three features. axis=0 removes the observation dimension, so the means describe features; the centered result retains both axes.

Run and check the result

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

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print("Shape:", x.shape)
print("Dtype:", x.dtype)
print("Devices:", x.devices())
print("Row sums:", x.sum(axis=1))
# Check tensor shape invariant: `x.shape == (2, 3)`
assert x.shape == (2, 3)
assert jnp.allclose(x.sum(axis=1), jnp.array([3., 12.]))

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

Run the 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: x has two observations and three features.
# Construct and reshape `x` into the target tensor dimensions.
x = jnp.arange(6, dtype=jnp.float32).reshape(2, 3)
# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
print("Shape:", x.shape)
print("Dtype:", x.dtype)
print("Devices:", x.devices())
print("Row sums:", x.sum(axis=1))
# Check tensor shape invariant: `x.shape == (2, 3)`
assert x.shape == (2, 3)
assert jnp.allclose(x.sum(axis=1), jnp.array([3., 12.]))

Shape: (2, 3)
Dtype: float32
Devices: {CpuDevice(id=0)}
Row sums: [ 3. 12.]
Shape: (2, 3)
Dtype: float32
Devices: {CpuDevice(id=0)}
Row sums: [ 3. 12.]
Means: [[1.5 2.5 3.5]]
Centered: [[-1.5 -1.5 -1.5]
 [ 1.5  1.5  1.5]]
Host dtype: float32
Expected element-count mismatch
PASS: welcome-02

Center each measurement

Predict before running: Predict the three column means and every centered value. Which mean is shared by the two observations?

# Experiment — Center each measurement: The calculation reuses one mean per column.
# Reduce along axis=0 to compute `means`.
means=x.mean(axis=0,keepdims=True)
centered=x-means
print("Means:",means)
print("Centered:",centered)
# Check tensor shape invariant: `means.shape==(1,3)`
assert means.shape==(1,3)
assert jnp.allclose(centered,jnp.array([[-1.5,-1.5,-1.5],[1.5,1.5,1.5]]))

The calculation reuses one mean per column. An explicit hand calculation verifies both the values and the meaning of the reduction.

Keep a host reference

Predict before running: Will copying to NumPy preserve values and shape? Does that copy describe device-resident computation?

# Experiment — Keep a host reference: The host reference is an inspection copy.
import numpy as np
# Convert `host` to a host NumPy array for inspection or verification.
host=np.asarray(x)
# Check tensor shape invariant: `host.shape==(2,3)`
assert host.shape==(2,3)
# Create evenly spaced index values in ``.
np.testing.assert_array_equal(host,np.arange(6).reshape(2,3))
print("Host dtype:",host.dtype)

The host reference is an inspection copy. Its existence says nothing about accelerator throughput.

Your exercise

Compute column sums instead. State the resulting shape before running the code.

# Exercise solution: Compute column sums instead.
# Reduce along axis=0 to compute `column_sums`.
column_sums = x.sum(axis=0)
# Check tensor shape invariant: `column_sums.shape == (3,)`
assert column_sums.shape == (3,)
assert jnp.allclose(column_sums, jnp.array([3., 5., 7.]))

Add a new observation

Practice

Append $[6,7,8]$. Predict row sums and column means, then check both.

Hint

Keep observation rows and measurement columns.

# Add a new observation (Practice): Changing the observation count changes axis 0, while the...
# Construct and reshape `table` into the target tensor dimensions.
table=jnp.arange(9,dtype=jnp.float32).reshape(3,3)
assert jnp.allclose(table.sum(axis=1),jnp.array([3.,12.,21.]))
assert jnp.allclose(table.mean(axis=0),jnp.array([3.,4.,5.]))

Diagnose an impossible reshape

Challenge

Try reshaping six values to $(2,4)$. Explain the failure and repair the intended three-measurement table.

Hint

Count required elements before changing the code.

# Diagnose an impossible reshape (Challenge): An incompatible reshape is an element-count error.
# Run the boundary check and catch the expected exception:
try:
    x.reshape(2,4)
except TypeError:
    print("Expected element-count mismatch")
else:
    raise AssertionError("Expected reshape to fail")
assert x.reshape(2,3).size==6

Checkpoint

For an array of shape $(2, 3)$, what shape results from sum(axis=0)?

1. $(2,)$
2. $(3,)$
3. $(2, 3)$

Diagnose the result

An incompatible reshape is an element-count error. Changing a dtype or moving devices cannot fix it. Choose a shape whose product matches six.

Carry forward

• The calculation reuses one mean per column. An explicit hand calculation verifies both the values and the meaning of the reduction.
• The host reference is an inspection copy. Its existence says nothing about accelerator throughput.

Your evidence

Keep the labeled table, row/column calculations, centered-data check, array dtype/device report and the failed reshape with its element-count diagnosis.

Primary references

• JAX arrays
• JAX data placement

Chapter 3 · welcome-cpu

Practice with four virtual CPU devices

Start four logical CPU devices, inspect row-sharded and replicated arrays, and diagnose an indivisible batch.

What you will be able to do

• Start four virtual CPU devices in a fresh process.
• Explain global shape, shard index and local shape.
• Compare row partitioning with replication.
• Diagnose divisibility and distinguish CPU practice from TPU evidence.

The problem

You can explore how JAX splits an array across devices using just your laptop. We’ll create four logical CPU devices, give each one a piece of an array, and compare the result with NumPy. Start with the setup and array/device lessons. These logical devices share your CPU: they help us learn placement, but they do not reproduce accelerator speed or memory.

The idea

Logical CPU devices let us practice how a global array is divided and addressed. They share a physical host. Treat them as a way to make placement visible, then use a real accelerator run to investigate accelerator memory, communication and speed.

Distinguish ownership from extra computing power

Suppose a global array has four rows and each logical device owns one row. The global shape still describes all four rows; a local shard describes the one row stored for that owner. Asking for the global shape and inspecting a local buffer are different operations.

The ownership heatmap assigns a category to each device. A darker device color does not mean more work or a faster device. Follow the row boundary and the device label together.

Set the logical-device configuration before JAX initializes its backend, then inspect the actual device count in a fresh process. Changing a flag after initialization does not retroactively repartition an already-running process.

Configure the runtime before it exists

Open the lesson workspace from setup and create main.py. The first block imports JAX, chooses CPU, and requests four CPU devices. Importing jax is allowed before the configuration; querying devices or creating an array can initialize the backend. That initialization is the boundary after which changing the count is too late.

Run python main.py as a fresh process. In Jupyter or Colab, restart the kernel/runtime and use Run all so this configuration is the first JAX operation. If you already ran an earlier lesson in that kernel, merely rerunning this cell cannot reliably rebuild the runtime. The explicit count check gives you a repair instruction instead of silently proceeding with one device. A terminal alternative is JAX_PLATFORMS=cpu JAX_NUM_CPU_DEVICES=4 python main.py; Windows PowerShell uses $env:JAX_PLATFORMS="cpu" and $env:JAX_NUM_CPU_DEVICES="4" before python main.py.

Distinguish logical devices from your CPU hardware

Four reported CPU devices do not mean you bought four CPUs or reserved four isolated cores. They share your computer and its resources. Each is a target JAX can place an array on. Keep the actual JAX version, device platform, process count and device count in your report so readers know what was tested.

This exercise has one controller process and four addressable devices. Multi-host JAX involves additional processes and coordination; this exercise does not test those. Even if you time an operation here, the observation belongs to this CPU setup and workload. It is distinct from whether the same partitioning is fast on a TPU.

A global array contains local pieces

The host array has shape $(4,2)$: four examples with two values each. Mesh assigns the four devices the named axis data. P("data",None) maps the first array dimension to that mesh axis and leaves the feature dimension unpartitioned. Each local shard therefore has shape $(1,2)$.

Use shard.index to identify its slice in the global array and shard.data to inspect its values. The global x.shape stays $(4,2)$; it does not become $(1,2)$ just because each device holds one row. Conversion with np.asarray(x) gathers the fully addressable global array to the host in this small single-process exercise.

Global rows [0,1] [2,3] [4,5] [6,7]
             ↓     ↓     ↓     ↓
CPU device   0     1     2     3
local shape (1,2) on every device; global shape (4,2)

Check arithmetic separately from placement

The function $2x+1$ is elementwise: a row does not need a value from a different row. For the first row $[0,1]$, predict $[1,3]$; for the final row $[6,7]$, predict $[13,15]$. The jit contract declares input and output placement explicitly. The NumPy check verifies all values, and shard inspections verify the intended partition. Neither alone verifies both.

block_until_ready waits for the result before reporting. It is useful when observing execution and essential for a future synchronized timing experiment, but no timings are claimed here. Four-way row partitioning requires a row count divisible by four. Try five rows deliberately, read the shape error, and choose a repair that preserves what your experiment means.

Four logical devices own four different rows

Predict before looking: Does each device hold the full array or one row?

Read the figure

Each cell represents a position in the global array, and its number identifies the logical CPU device that owns that position. The colors are categories: device $3$ is not “more powerful” than device $0$.

Both cells in row $0$ belong to device $0$; both cells in row $1$ belong to device $1$, and the same pattern continues through row $3$. Horizontal bands show that this layout partitions rows while keeping each row’s features together.

Connect it to the computation

The global shape is $(4,2)$, so splitting its first axis across four devices gives each device a local piece with shape $(1,2)$. This is the spatial meaning of the row partition in the code. The numbers in the figure are ownership labels, not the original array’s contents.

You can check the interpretation by locating row $2$, feature $1$: its owner is device $2$. These logical devices share the host CPU. The figure establishes where values are placed, not a measured speedup or the behavior of four physical accelerators.

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

Start a fresh CPU runtime

Create main.py in your course workspace. Paste this block first. If using a notebook, restart its kernel before running any cell.

# Step 1 — Start a fresh CPU runtime: Configuration happens before devices() or array creation...
import jax
# Run this file in a fresh process; in a notebook restart the kernel first.
jax.config.update("jax_platforms", "cpu")
# Update state in place with the new values.
jax.config.update("jax_num_cpu_devices", 4)
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
# Query the active JAX devices into `devices`.
devices = jax.devices("cpu")
# Guard input contract (`len(devices) != 4`) and fail fast if violated.
if len(devices) != 4:
    raise RuntimeError("Expected four CPU devices. Restart the notebook kernel, then Run all; do not run an earlier JAX cell first.")
# Configure multi-device placement / sharding specification (`mesh`).
mesh = Mesh(np.array(devices), ("data",))
# Configure multi-device placement / sharding specification (`rows`).
rows = NamedSharding(mesh, P("data", None))
# Configure multi-device placement / sharding specification (`replicated`).
replicated = NamedSharding(mesh, P())

Configuration happens before devices() or array creation initializes a backend. This file explicitly chooses CPU even on a GPU machine.

Place and compute

Append this block in the same file. Predict the shapes and values before running.

# Step 2 — Place and compute: Mesh axis names describe placement.
# Construct and reshape `host` into the target tensor dimensions.
host = np.arange(8, dtype=np.float32).reshape(4, 2)
# Place `x` explicitly onto the target JAX device.
x = jax.device_put(host, rows)
# Wrap with `jax.jit` (`y`) so XLA traces and compiles the function.
y = jax.jit(lambda a: 2 * a + 1, in_shardings=rows, out_shardings=rows)(x)
# Synchronize host execution until asynchronous device computation completes.
y.block_until_ready()

Mesh axis names describe placement. They are separate from the numerical array dimensions.

Inspect and verify

Append the checks, save the file and run python main.py with the setup lesson environment.

# Step 3 — Inspect and verify: Assertions compare against host calculations or hand-derived...
print("JAX version:", jax.__version__)
print("CPU devices:", len(devices))
for shard in x.addressable_shards:
    print("Device", shard.device.id, "index", shard.index, "values", np.asarray(shard.data))
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(y), 2 * host + 1)
assert len(x.addressable_shards) == 4
# Check tensor shape invariant: `all(shard.data.shape == (1, 2) for shard in x.addressable_shards)`
assert all(shard.data.shape == (1, 2) for shard in x.addressable_shards)
print("Global shape:", x.shape, "result:", np.asarray(y).tolist())

Assertions compare against host calculations or hand-derived values; printing a sharding object alone is separate from correctness.

Run the example

# Step 1 — Start a fresh CPU runtime: Configuration happens before devices() or array creation...
import jax
# Run this file in a fresh process; in a notebook restart the kernel first.
jax.config.update("jax_platforms", "cpu")
# Update state in place with the new values.
jax.config.update("jax_num_cpu_devices", 4)
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
# Query the active JAX devices into `devices`.
devices = jax.devices("cpu")
# Guard input contract (`len(devices) != 4`) and fail fast if violated.
if len(devices) != 4:
    raise RuntimeError("Expected four CPU devices. Restart the notebook kernel, then Run all; do not run an earlier JAX cell first.")
# Configure multi-device placement / sharding specification (`mesh`).
mesh = Mesh(np.array(devices), ("data",))
# Configure multi-device placement / sharding specification (`rows`).
rows = NamedSharding(mesh, P("data", None))
# Configure multi-device placement / sharding specification (`replicated`).
replicated = NamedSharding(mesh, P())

# Step 2 — Place and compute: Mesh axis names describe placement.
# Construct and reshape `host` into the target tensor dimensions.
host = np.arange(8, dtype=np.float32).reshape(4, 2)
# Place `x` explicitly onto the target JAX device.
x = jax.device_put(host, rows)
# Wrap with `jax.jit` (`y`) so XLA traces and compiles the function.
y = jax.jit(lambda a: 2 * a + 1, in_shardings=rows, out_shardings=rows)(x)
# Synchronize host execution until asynchronous device computation completes.
y.block_until_ready()

# Step 3 — Inspect and verify: Assertions compare against host calculations or hand-derived...
print("JAX version:", jax.__version__)
print("CPU devices:", len(devices))
for shard in x.addressable_shards:
    print("Device", shard.device.id, "index", shard.index, "values", np.asarray(shard.data))
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(y), 2 * host + 1)
assert len(x.addressable_shards) == 4
# Check tensor shape invariant: `all(shard.data.shape == (1, 2) for shard in x.addressable_shards)`
assert all(shard.data.shape == (1, 2) for shard in x.addressable_shards)
print("Global shape:", x.shape, "result:", np.asarray(y).tolist())

JAX version: 0.9.2
CPU devices: 4
Device 0 index (slice(0, 1, None), slice(None, None, None)) values [[0. 1.]]
Device 1 index (slice(1, 2, None), slice(None, None, None)) values [[2. 3.]]
Device 2 index (slice(2, 3, None), slice(None, None, None)) values [[4. 5.]]
Device 3 index (slice(3, 4, None), slice(None, None, None)) values [[6. 7.]]
Global shape: (4, 2) result: [[1.0, 3.0], [5.0, 7.0], [9.0, 11.0], [13.0, 15.0]]
JAX version: 0.9.2
CPU devices: 4
Device 0 index (slice(0, 1, None), slice(None, None, None)) values [[0. 1.]]
Device 1 index (slice(1, 2, None), slice(None, None, None)) values [[2. 3.]]
Device 2 index (slice(2, 3, None), slice(None, None, None)) values [[4. 5.]]
Device 3 index (slice(3, 4, None), slice(None, None, None)) values [[6. 7.]]
Global shape: (4, 2) result: [[1.0, 3.0], [5.0, 7.0], [9.0, 11.0], [13.0, 15.0]]
Negative input result: [[-3.0, -5.0], [-7.0, -9.0], [-11.0, -13.0], [-15.0, -17.0]]
Replicated local shapes: [(4, 2), (4, 2), (4, 2), (4, 2)]
Expected indivisible row dimension
Expected indivisible row dimension
PASS: welcome-cpu

Keep the shape, change the values

Predict before running: If every input value becomes negative, do the number of devices or local shard shapes change?

# Experiment — Keep the shape, change the values: Placement follows the declared shape/specification, not the...
negative = -host - 2
# Place `negative_x` explicitly onto the target JAX device.
negative_x = jax.device_put(negative, rows)
# Wrap with `jax.jit` (`negative_y`) so XLA traces and compiles the function.
negative_y = jax.jit(lambda a: 2*a+1, in_shardings=rows, out_shardings=rows)(negative_x)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(negative_y), 2*negative+1)
# Check tensor shape invariant: `all(s.data.shape == (1,2) for s in negative_x.addressable_shards)`
assert all(s.data.shape == (1,2) for s in negative_x.addressable_shards)
print("Negative input result:", np.asarray(negative_y).tolist())

Placement follows the declared shape/specification, not the signs or magnitudes of array elements.

Replicate instead of partitioning

Predict before running: With P(), which values should each CPU device receive?

# Experiment — Replicate instead of partitioning: Replication stores full copies.
copies = jax.device_put(host, replicated)
assert copies.is_fully_replicated
assert len(copies.addressable_shards) == 4
for shard in copies.addressable_shards:
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_array_equal(np.asarray(shard.data), host)
print("Replicated local shapes:", [s.data.shape for s in copies.addressable_shards])

Replication stores full copies. A global array shape alone cannot tell you how much of it lives on each device.

Your exercise

Predict and verify the per-device shapes for twelve rows, then demonstrate how replication repairs an indivisible five-row placement without deleting data.

# Exercise solution: Predict and verify the per-device shapes for twelve rows, then...
# Construct and reshape `twelve` into the target tensor dimensions.
twelve = np.arange(24,dtype=np.float32).reshape(12,2)
# Place `z` explicitly onto the target JAX device.
z = jax.device_put(twelve, rows)
# Wrap with `jax.jit` (`z2`) so XLA traces and compiles the function.
z2 = jax.jit(lambda a:a*a, in_shardings=rows, out_shardings=rows)(z)
# Check tensor shape invariant: `all(s.data.shape == (3,2) for s in z2.addressable_shards)`
assert all(s.data.shape == (3,2) for s in z2.addressable_shards)
for shard in z2.addressable_shards:
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_array_equal(np.asarray(shard.data), (twelve*twelve)[shard.index])

# Construct and reshape `uneven` into the target tensor dimensions.
uneven = np.arange(10,dtype=np.float32).reshape(5,2)
# Run the boundary check and catch the expected exception:
try:
    jax.device_put(uneven, rows)
except ValueError:
    print("Expected indivisible row dimension")
else:
    raise AssertionError("Expected divisibility failure")
# Place `fixed` explicitly onto the target JAX device.
fixed = jax.device_put(uneven, replicated)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(fixed), uneven)
assert fixed.is_fully_replicated

Move from four rows to twelve

Challenge

Create twelve two-feature rows with values $0$ through $23$. Predict shard shape and check every slice of a squared output.

Hint

Use each shard index to select the corresponding NumPy reference slice.

# Move from four rows to twelve (Challenge): The twelve rows split into four groups of three.
# Construct and reshape `twelve` into the target tensor dimensions.
twelve = np.arange(24,dtype=np.float32).reshape(12,2)
# Place `z` explicitly onto the target JAX device.
z = jax.device_put(twelve, rows)
# Wrap with `jax.jit` (`z2`) so XLA traces and compiles the function.
z2 = jax.jit(lambda a:a*a, in_shardings=rows, out_shardings=rows)(z)
# Check tensor shape invariant: `all(s.data.shape == (3,2) for s in z2.addressable_shards)`
assert all(s.data.shape == (3,2) for s in z2.addressable_shards)
for shard in z2.addressable_shards:
    # Convert `` to a host NumPy array for inspection or verification.
    np.testing.assert_array_equal(np.asarray(shard.data), (twelve*twelve)[shard.index])

Repair a batch that does not divide evenly

Challenge

Place a $(5,2)$ array with the same row sharding. Capture the failure. Repair this placement demonstration by replicating it, and explain why that does not prove the partitioned algorithm supports uneven batches.

Hint

Do not drop a row to hide the error.

# Repair a batch that does not divide evenly (Challenge): Five rows cannot be evenly split across the four-way data axis.
# Construct and reshape `uneven` into the target tensor dimensions.
uneven = np.arange(10,dtype=np.float32).reshape(5,2)
# Run the boundary check and catch the expected exception:
try:
    jax.device_put(uneven, rows)
except ValueError:
    print("Expected indivisible row dimension")
else:
    raise AssertionError("Expected divisibility failure")
# Place `fixed` explicitly onto the target JAX device.
fixed = jax.device_put(uneven, replicated)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(fixed), uneven)
assert fixed.is_fully_replicated

Checkpoint

What does a $(12,2)$ global array sharded with P("data",None) over four devices mean?

1. Each device holds a $(3,2)$ slice and the global array remains $(12,2)$.
2. Each device holds the complete array.
3. Each device represents a separate TPU chip.

Diagnose the result

Five rows cannot be evenly split across the four-way data axis. Replication preserves all observations but changes the memory/communication plan; padding requires a mask when computing means.

Carry forward

• Configure the runtime before it exists
• Distinguish logical devices from your CPU hardware
• A global array contains local pieces
• Check arithmetic separately from placement

Your evidence

Keep package/device reports, global and local shapes, shard index/value tables, NumPy comparisons, changed-shape predictions, and the five-row failure/repair. State that the run uses one host and does not validate TPU performance.

Primary references

• JAX CPU device configuration (set before initialization)
• Distributed arrays and automatic parallelization

Chapter 4 · welcome-03

Move your experiment to a TPU

Provision a hosted TPU runtime or single-host Cloud TPU VM, launch an explicit TPU target path over SSH that fails when TPU is unavailable, verify the device of a completed prediction against NumPy, and set up parallel TPU execution for the rest of the course.

What you will be able to do

• Verify the device of a completed prediction and its independent numerical reference.
• Provision a hosted TPU notebook or single-host Google Cloud TPU VM and verify device discovery.
• Launch an explicit TPU target path over SSH that fails when TPU is unavailable, then clean up the resource.
• Run subsequent course lessons and companion labs on TPU in parallel with local CPU practice.

The problem

Your notebook imports JAX successfully, but is the prediction actually running on a TPU? We will build a portable device check with an independent numerical answer, then require the requested backend so an unavailable accelerator cannot silently become a successful CPU demonstration.

The idea

Moving an experiment to a TPU gives us two questions: did it calculate the intended result, and did it actually execute on the requested device? We answer the first with an independent numerical reference and the second with an observed device receipt.

Check arithmetic and placement separately

Consider a prediction that adds a scalar bias after a dot product. Reducing the bias by $0.25$ should reduce every prediction by $0.25$, whatever compatible device performs the computation. That is a numerical invariant you can derive before running.

Now inspect where the completed result resides. A correct NumPy comparison on CPU does not show that a TPU was selected. Conversely, a TPU device label does not show that the formula or input data was correct. Keep both checks.

Wait for the result before recording completion or timing. A returned array handle can precede finished device work. Preserve the requested platform, observed platform, environment and error tolerance alongside the values.

Carry a prediction, not just an installation

We will move one small function between runtimes: a matrix-vector prediction with a scalar bias. Each row is an observation; three columns are three features. The weights are fixed so we can calculate the answer independently. For the first row, $[0,1/8,2/8]$, the prediction is $0-1/32+1/4+1/8=11/32$. A working package import is useful, but it is separate from where this computation finished.

Establish the CPU reference first

Save the assembled code as main.py in your course workspace and run it in a fresh process. By default COURSE_EXPECT_PLATFORM is cpu. The program selects that backend explicitly, places both inputs on its first device and waits for the compiled prediction to finish. NumPy independently calculates the same arithmetic. Keep the environment report and predictions.

The tolerance allows small floating-point differences; a successful comparison means this prediction agrees within that tolerance. It is separate from accuracy for a larger model. We use float32 inputs deliberately. Implicit dtype changes and reduced-precision matrix multiplication should be separate experiments after the placement check.

Run the CPU reference

# Run run the cpu reference using the course Python environment
python main.py

Expected:

The report names cpu, shape [8], float32, a completed output device and a small maximum absolute error. The first prediction is 0.34375.

Get a TPU running: hosted notebook or Google Cloud TPU VM

You do not need to wait until the end of the course to use a TPU. There are two practical ways to get a TPU running right now so you can run experiments alongside the course:

1. Hosted notebook TPU (Google Colab or Kaggle): open Runtime -> Change runtime type, select a TPU accelerator (such as TPU v5e-1 or v6e-1), and restart the kernel before importing JAX so the process does not lock in a CPU backend.
2. Google Cloud single-host TPU VM: install the Google Cloud CLI on your laptop, authenticate with gcloud auth login, and create a single-host TPU VM (such as v5litepod-1, v5litepod-4 or v6e-1) in an approved project and zone. Wait until the VM state is READY. Remember that a Cloud TPU VM is a billable machine whose lifetime is separate from your Python process: when Python exits, the VM keeps running until you delete it or its bounded max-run-duration expires.

From your laptop: provision a single-host Cloud TPU VM and check READY state

# From your laptop: provision a single-host Cloud TPU VM and check READY state
export PROJECT_ID="REPLACE_WITH_YOUR_PROJECT"
export ZONE="us-central1-a"
export TPU_NAME="jax-course-tpu-01"
gcloud compute tpus tpu-vm create "$TPU_NAME" \
  --project="$PROJECT_ID" --zone="$ZONE" \
  --accelerator-type="v5litepod-1" --version="v2-alpha-tpuv5-lite"
gcloud compute tpus tpu-vm describe "$TPU_NAME" \
  --project="$PROJECT_ID" --zone="$ZONE" --format="value(state)"

Expected:

The Cloud TPU VM is created in your project/zone and describe prints READY. For bounded Flex-start templates or queued resources, see the TPU and Google Cloud practical guide.

Inside the TPU VM or hosted TPU runtime: install JAX for TPU and verify devices

# Run inside the tpu vm or hosted tpu runtime: install jax for tpu and verify devices using the course Python environment
python3 -m venv .venv
.venv/bin/python -m pip install --upgrade "jax[tpu]" numpy
.venv/bin/python -c "import jax; print(jax.default_backend(), jax.devices('tpu'))"

Expected:

The TPU environment installs compatible JAX and libtpu packages and prints tpu with the visible TpuDevice list. Record the resolved versions; installation alone is not the execution check.

Launch the experiment on TPU, retrieve the receipt, and delete the VM

Once your TPU environment is ready, copy main.py to the TPU host, launch it in a fresh process with COURSE_EXPECT_PLATFORM=tpu and JAX_PLATFORMS=tpu, and copy the output report back to your laptop. The call to jax.devices(expected) requests TPU explicitly and fails immediately if the TPU backend is unavailable.

This is a single-controller, single-device placement check even if the runtime exposes four or eight chips. Multi-chip sharding on one host and multi-host launches with jax.distributed.initialize() are covered in the virtual-CPU and distributed lessons. As soon as you copy your evidence back to your laptop, delete the Cloud TPU VM and verify that the instance list is empty.

Require TPU execution inside an active TPU session

# Require TPU execution inside an active TPU session
COURSE_EXPECT_PLATFORM=tpu JAX_PLATFORMS=tpu python main.py

Expected:

On a functioning TPU runtime the report names tpu and all assertions pass. An unavailable TPU must raise an error instead of producing a CPU success report.

From your laptop: copy main.py, launch over SSH, copy back evidence, and delete the TPU VM

# From your laptop: copy main.py, launch over SSH, copy back evidence, and delete the TPU VM
gcloud compute tpus tpu-vm scp main.py "$TPU_NAME:~/main.py" --project="$PROJECT_ID" --zone="$ZONE"
gcloud compute tpus tpu-vm ssh "$TPU_NAME" --project="$PROJECT_ID" --zone="$ZONE" \
  --command="COURSE_EXPECT_PLATFORM=tpu JAX_PLATFORMS=tpu .venv/bin/python ~/main.py | tee ~/tpu-report.txt"
gcloud compute tpus tpu-vm scp "$TPU_NAME:~/tpu-report.txt" ./tpu-report.txt --project="$PROJECT_ID" --zone="$ZONE"
gcloud compute tpus tpu-vm delete "$TPU_NAME" --project="$PROJECT_ID" --zone="$ZONE" --quiet
gcloud compute tpus tpu-vm list --project="$PROJECT_ID" --zone="$ZONE"

Expected:

The remote command runs on the TPU VM, writes tpu-report.txt locally, deletes the TPU VM, and confirms no running TPU VM remains.

Run the rest of the course on TPU in parallel

Use this same launch pattern in parallel with every phase that follows. Every authored lesson in the course provides a standalone script under exercises/<lesson-id>.py (and a notebook under exercises/<lesson-id>.ipynb). As you work through Arrays (Phase 01), Transforms (Phase 02), State (Phase 03), Optimization (Phase 04), Networks (Phase 05) and Recovery (Phase 06), run the CPU reference locally first, then upload the workspace or script to your TPU runtime and run it with JAX_PLATFORMS=tpu.

In parallel with the early phases, work through the two standalone TPU guides included in the course workspace (jax-tpu-gcp.zip): use resources/tpu-gcp/launch.py to practice supervised TPU job lifecycles and checkpoint recovery, and use resources/tpu-gcp/precision_profile.py to measure BF16/INT8 numerical error and capture XProf device traces before you scale up to full Transformers and distributed training.

Run any course lesson script and the companion TPU labs on a TPU host

# Run any course lesson script and the companion TPU labs on a TPU host
JAX_PLATFORMS=tpu .venv/bin/python exercises/arrays-01.py
.venv/bin/python resources/tpu-gcp/launch.py --platform tpu --run-dir runs/tpu-first --steps 8 --timeout 180
.venv/bin/python resources/tpu-gcp/precision_profile.py --platform tpu --run-dir runs/precision-tpu --steps 20 --trace

Expected:

The lesson script and companion TPU labs run on the TPU backend, writing verifiable device receipts, checkpoints, numerical error reports and XProf trace files.

Diagnose the boundary that failed

If import fails, check the interpreter and installation. If explicit backend selection fails, check whether you are in the intended runtime and whether its TPU service is available. If device placement succeeds but the independent comparison fails, inspect shapes, dtype, precision and the arithmetic. Do not widen tolerance just to hide a large error.

An old notebook kernel may retain an initialized CPU backend or stale variables. Restart it, set runtime choices before the first device operation and run every cell. A CPU plot from the recorded course run remains a CPU plot when viewed on a machine that happens to have TPU access; only a fresh recorded execution supplies new evidence.

The same prediction checked two ways

Predict before looking: Should the JAX line overlap the independent NumPy points, and would overlap alone prove TPU execution?

Read the figure

The horizontal axis is the observation row index, starting at $0$. The vertical axis is the scalar prediction. The two series use the same inputs and weights: one comes from the completed JAX array, the other from NumPy. They overlap in this recorded CPU example. The first value is $11/32$; the constant rise of $15/32$ follows from increasing every feature by $3/8$ while the weights sum to $5/4$.

Connect it to the computation

The line is a check of arithmetic and ordering for these eight observations. Read the printed platform report to determine placement; a curve cannot identify the device that produced it. The changed-bias exercise translates this line downward, while row reversal changes the order. This recorded figure comes from the CPU path and carries no TPU timing or correctness claim.

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

Choose the requested platform

Add this block to main.py after the preceding block. Run the assembled file in a fresh process.

# Step 1 — Choose the requested platform: Keep the requested backend, observed output placement and...
import os
import json
import platform
import numpy as np
import jax
import jax.numpy as jnp

Keep the requested backend, observed output placement and independent arithmetic check together in the final report.

Place inputs and compare a completed prediction

Add this block to main.py after the preceding block. Run the assembled file in a fresh process.

# Step 2 — Place inputs and compare a completed prediction: Keep the requested backend, observed output placement and...
expected = os.environ.get('COURSE_EXPECT_PLATFORM', 'cpu')
# Guard input contract (`expected not in {'cpu', 'tpu'}`) and fail fast if violated.
if expected not in {'cpu', 'tpu'}:
    raise ValueError('COURSE_EXPECT_PLATFORM must be cpu or tpu')
# Selecting a requested backend must fail when it is unavailable.
devices = jax.devices(expected)
assert devices and all(d.platform == expected for d in devices)
# Construct and reshape `x_host` into the target tensor dimensions.
x_host = np.arange(24, dtype=np.float32).reshape(8, 3) / 8
w_host = np.array([0.5, -0.25, 1.0], dtype=np.float32)
# Place `x` explicitly onto the target JAX device.
x = jax.device_put(x_host, devices[0])
# Place `w` explicitly onto the target JAX device.
w = jax.device_put(w_host, devices[0])
# Wrap with `jax.jit` (`predict`) so XLA traces and compiles the function.
predict = jax.jit(lambda a, b: a @ b + jnp.float32(0.125))
y = predict(x, w)
# Synchronize host execution until asynchronous device computation completes.
y.block_until_ready()
reference = x_host @ w_host + np.float32(0.125)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(y), reference, rtol=1e-5, atol=1e-5)
assert {d.platform for d in y.devices()} == {expected}
# Convert `report` to a host NumPy array for inspection or verification.
report = dict(expected=expected, output_devices=[str(d) for d in y.devices()],
              platform=next(iter(y.devices())).platform, jax=jax.__version__,
              python=platform.python_version(), shape=list(y.shape), dtype=str(y.dtype),
              max_absolute_error=float(np.max(np.abs(np.asarray(y)-reference))),
              process_index=jax.process_index(), process_count=jax.process_count())
print(json.dumps(report, indent=2))
print('Completed prediction:', np.asarray(y).tolist())

Keep the requested backend, observed output placement and independent arithmetic check together in the final report.

Step 3: Verify invariants on the completed state

Run the final shape and numerical assertions to confirm the state built in Steps 1 and 2.

# Assert that the device output matches the NumPy reference and print the portability report:
assert report["max_absolute_error"] < 1e-6
print(json.dumps(report, indent=2))
print("Completed prediction:", np.asarray(y).tolist())

Checking these invariants confirms the computation is ready for the full worked experiment.

Run the example

# Step 1 — Choose the requested platform: Keep the requested backend, observed output placement and...
import os
import json
import platform
import numpy as np
import jax
import jax.numpy as jnp

# Step 2 — Place inputs and compare a completed prediction: Keep the requested backend, observed output placement and...
expected = os.environ.get('COURSE_EXPECT_PLATFORM', 'cpu')
# Guard input contract (`expected not in {'cpu', 'tpu'}`) and fail fast if violated.
if expected not in {'cpu', 'tpu'}:
    raise ValueError('COURSE_EXPECT_PLATFORM must be cpu or tpu')
# Selecting a requested backend must fail when it is unavailable.
devices = jax.devices(expected)
assert devices and all(d.platform == expected for d in devices)
# Construct and reshape `x_host` into the target tensor dimensions.
x_host = np.arange(24, dtype=np.float32).reshape(8, 3) / 8
w_host = np.array([0.5, -0.25, 1.0], dtype=np.float32)
# Place `x` explicitly onto the target JAX device.
x = jax.device_put(x_host, devices[0])
# Place `w` explicitly onto the target JAX device.
w = jax.device_put(w_host, devices[0])
# Wrap with `jax.jit` (`predict`) so XLA traces and compiles the function.
predict = jax.jit(lambda a, b: a @ b + jnp.float32(0.125))
y = predict(x, w)
# Synchronize host execution until asynchronous device computation completes.
y.block_until_ready()
reference = x_host @ w_host + np.float32(0.125)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(y), reference, rtol=1e-5, atol=1e-5)
assert {d.platform for d in y.devices()} == {expected}
# Convert `report` to a host NumPy array for inspection or verification.
report = dict(expected=expected, output_devices=[str(d) for d in y.devices()],
              platform=next(iter(y.devices())).platform, jax=jax.__version__,
              python=platform.python_version(), shape=list(y.shape), dtype=str(y.dtype),
              max_absolute_error=float(np.max(np.abs(np.asarray(y)-reference))),
              process_index=jax.process_index(), process_count=jax.process_count())
print(json.dumps(report, indent=2))
print('Completed prediction:', np.asarray(y).tolist())

{
  "expected": "cpu",
  "output_devices": [
    "TFRT_CPU_0"
  ],
  "platform": "cpu",
  "jax": "0.9.2",
  "python": "3.14.3",
  "shape": [
    8
  ],
  "dtype": "float32",
  "max_absolute_error": 0.0,
  "process_index": 0,
  "process_count": 1
}
Completed prediction: [0.34375, 0.8125, 1.28125, 1.75, 2.21875, 2.6875, 3.15625, 3.625]
{
  "expected": "cpu",
  "output_devices": [
    "TFRT_CPU_0"
  ],
  "platform": "cpu",
  "jax": "0.9.2",
  "python": "3.14.3",
  "shape": [
    8
  ],
  "dtype": "float32",
  "max_absolute_error": 0.0,
  "process_index": 0,
  "process_count": 1
}
Completed prediction: [0.34375, 0.8125, 1.28125, 1.75, 2.21875, 2.6875, 3.15625, 3.625]
{
  "expected": "cpu",
  "output_devices": [
    "TFRT_CPU_0"
  ],
  "platform": "cpu",
  "jax": "0.9.2",
  "python": "3.14.3",
  "shape": [
    8
  ],
  "dtype": "float32",
  "max_absolute_error": 0.0,
  "process_index": 0,
  "process_count": 1
}
Completed prediction: [0.34375, 0.8125, 1.28125, 1.75, 2.21875, 2.6875, 3.15625, 3.625]
Mismatched platform report rejected.
Changed bias, first prediction: -0.03125
Row reversal preserved values and requested placement.
PASS: welcome-03

Reject a mismatched report

Predict before running: Would a correct numeric answer be enough if the report says CPU but the requested target is TPU?

# Experiment — Reject a mismatched report: Numerical agreement and requested placement are separate...
def verify_report(record, required):
    # Guard input contract (`record['platform'] != required or record['expected'] != required`) and fail fast if violated.
    if record['platform'] != required or record['expected'] != required:
        raise ValueError('requested and observed platform disagree')
    # Guard input contract (`record['shape'] != [8] or record['dtype'] != 'float32'`) and fail fast if violated.
    if record['shape'] != [8] or record['dtype'] != 'float32':
        raise ValueError('prediction contract changed')
    return True
assert verify_report(report, expected)
wrong = {**report, 'platform': 'cpu' if expected == 'tpu' else 'tpu'}
# Run the boundary check and catch the expected exception:
try:
    verify_report(wrong, expected)
except ValueError:
    print('Mismatched platform report rejected.')
else:
    raise AssertionError('a mismatched report was accepted')

Numerical agreement and requested placement are separate requirements. This experiment tests report validation; the main program performs the actual backend selection and array-placement check.

Your exercise

Predict the effect of replacing the bias by minus one quarter. Run a second compiled prediction on the selected device. Compare against the original prediction plus the expected constant shift and an independent NumPy calculation.

# Exercise solution: Predict the effect of replacing the bias by minus one quarter.
# Wrap with `jax.jit` (`changed_predict`) so XLA traces and compiles the function.
changed_predict = jax.jit(lambda a, b: a @ b - jnp.float32(0.25))
changed = changed_predict(x, w)
# Synchronize host execution until asynchronous device computation completes.
changed.block_until_ready()
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(changed), reference - 0.375, atol=1e-5, rtol=1e-5)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(changed), x_host @ w_host - 0.25, atol=1e-5, rtol=1e-5)
assert {d.platform for d in changed.devices()} == {expected}
print('Changed bias, first prediction:', float(changed[0]))

Test row order without changing the model

Transfer

Reverse the observation rows and predict the result order before execution. Verify the first and last values, the complete vector and device placement. Explain why this is a stronger check than only comparing the vector sum.

Hint

A row-wise affine prediction commutes with row permutation. A sum would conceal an ordering bug.

# Test row order without changing the model (Transfer): Reversing rows should reverse predictions because each...
reversed_x = jax.device_put(x_host[::-1].copy(), devices[0])
reversed_y = predict(reversed_x, w)
# Synchronize host execution until asynchronous device computation completes.
reversed_y.block_until_ready()
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(reversed_y), reference[::-1], atol=1e-5, rtol=1e-5)
assert {d.platform for d in reversed_y.devices()} == {expected}
assert np.isclose(float(reversed_y[0]), 3.625)
assert np.isclose(float(reversed_y[-1]), 0.34375)
print('Row reversal preserved values and requested placement.')

Checkpoint

The values match NumPy, but the output array is on CPU when the requested platform was TPU. What has been demonstrated?

1. The TPU installation and its execution are validated.
2. The arithmetic matches the reference on the fallback CPU backend, so the requested TPU backend was not active.
3. The model must be numerically wrong.

Diagnose the result

Read the first failing boundary: interpreter/import, explicit backend discovery, output placement or numerical comparison. Keep the original error and requested platform in the report. Restart a stale notebook kernel before trying the target again.

Carry forward

• Provision a hosted TPU runtime or single-host Cloud TPU VM early and run course experiments in parallel with your CPU baseline.
• Check an actual completed array, not only the installed package.
• Request TPU explicitly so its absence produces an error, and delete Cloud TPU VMs after copying your evidence.

Your evidence

Keep separate CPU and TPU reports with requested/observed platform, package versions, completed predictions, NumPy error and changed-bias and row-permutation checks. Until the target command succeeds, TPU execution remains unverified.

Primary references

• JAX installation and TPU support
• Manage Cloud TPU VMs with gcloud
• Course guide: From a laptop experiment to a TPU training run
• Course guide: TPU generations, precision and profiling
• JAX device placement

Chapter 5 · arrays-01

From NumPy to jax.numpy

Build feature-wise preprocessing with fitted statistics and verify behavior on new and constant data.

What you will be able to do

• Translate a feature-wise NumPy calculation to jax.numpy.
• Predict the shapes and values of reductions and broadcasted centering.
• Fit preprocessing statistics on training data and reuse them on new data.
• Handle constant features explicitly and compare against a NumPy reference.

The problem

Suppose your dataset records two measurements on very different scales. How can we put them on a more comparable footing before training a model? We’ll center each feature and divide by its scale. First we’ll work through three rows by hand; then we’ll write the JAX version and try new data, including a feature that never changes.

The idea

Feature preprocessing answers a specific question: how does each measurement compare with the training data's scale and center? Compute statistics along the observation axis, then reuse those fitted statistics. A held-out example must not help choose the transformation used to evaluate it.

Why held-out data should not center itself

Use a small new case: training values for one feature are $2$ and $6$, so their mean is $4$. Centered training values become $-2$ and $2$. A held-out value of $10$ becomes $6$ when we apply the training mean.

If we center that single held-out value using its own mean, it becomes zero. The computation succeeds, but it has erased the evidence that the new value lies far from the training center. With several features, retain one fitted mean per feature rather than collapsing the entire table to one mean.

The lesson's centering plot compares values before and after the same columnwise operation. It explains a translation of each feature, not a guarantee that future observations will have zero mean. Add scaling only after specifying the variance convention and how constant features are handled.

Follow an axis through a reduction

A shape $(3, 2)$ means three observations with two feature values each. mean(axis=0) reduces the three observations and leaves two feature means, so its output shape is $(2,)$. mean(axis=1) instead averages unlike feature values within each observation, leaving shape $(3,)$. mean() with no axis reduces every value to a scalar.

All three calculations are legal. Only the first answers the question about each feature across observations. A valid shape is not evidence of a meaningful statistic: if the features represent different units, averaging them within a row might have no useful interpretation.

(N,D) -- mean over axis 0 --> (D,)
(N,D) -- mean over axis 1 --> (N,)
(N,D) -- mean over all axes --> ()

Derive the centering and scale values

Think of one column at a time. The symbol $x_{ij}$ means the value in row $i$, feature $j$, and $N$ is the number of training rows. The mean $\mu_j$ is the column’s center. The variance $s_j^2$ averages squared distances from that center; its square root $s_j$ is the scale we divide by.

For the first feature $[1, 3, 5]$, the mean is $3$. The centered values are $[-2, 0, 2]$, so the variance is $8/3$. The second feature $[10, 14, 18]$ centers to $[-4, 0, 4]$ and has variance $32/3$. Its scale is twice as large, so both columns have the same standardized pattern.

A constant column has scale zero. We use a safe denominator of $1$ for that column, which leaves its centered training values at zero. Learn the mean and scale from training data, then reuse them on new data; recomputing them on each incoming batch changes the meaning of a feature.

Fit statistics once; apply them consistently

Centering training data using its own mean is useful. Recomputing the mean separately for each validation batch changes the transformation and incorporates information from that batch. A model then sees features represented in different coordinate systems.

Separate fit_standardizer(training_data), which computes statistics, from apply_standardizer(data, statistics), which reuses them. The same rule will apply to training/validation splits, inference requests, and checkpointed preprocessing state. A new batch is not expected to have zero mean after applying training statistics; that difference can be meaningful.

Give constant features an explicit policy

A constant feature has zero standard deviation. Dividing its zero centered values by zero produces an undefined result. Adding epsilon to every scale is a simple bounded choice, but it perturbs the scale of every feature and can amplify a tiny feature with near-zero variability.

For this lesson, use $s=1$ for exactly constant columns and their training standard deviation otherwise. Training values in a constant column then transform to zero without division by zero. A changed value in new data produces a nonzero centered result. This policy is explicit and testable; it is not a universal prescription for every dataset.

Use NumPy as a numerical reference, not an execution substitute

Compare a tiny JAX calculation with NumPy using explicitly matching float32 data and the same reduction definitions. Transfer and conversion are fine for this test. Converting an intermediate to NumPy inside a transformed differentiable computation is a different operation and can break its intended transform behavior.

The output of a jax.numpy operation is a JAX array with its own device behavior. We will explore devices and compilation elsewhere. At this point, focus on matching values, preserving shapes, and keeping preprocessing semantics explicit.

Center each feature around its own mean

Predict before looking: Does centering make every entry zero, or only each column mean?

Read the figure

The horizontal axis indexes the three observations; the vertical axis gives feature values. Compare each original feature with its centered version in the legend. The connecting lines help you follow corresponding entries; they are not a time series.

Feature $0$ changes from $(1,3,5)$ to $(-2,0,2)$, a downward shift of $3$. Feature $1$ changes from $(10,14,18)$ to $(-4,0,4)$, a downward shift of $14$. Both centered series pass through zero at the middle observation.

Connect it to the computation

Those shifts are the separate column means. Subtracting a mean moves a whole feature without changing the gaps between its observations: feature $0$ still increases by $2$, while feature $1$ still increases by $4$. That is why each centered line stays parallel to its original line.

Centering makes each column’s mean zero, but does not give the columns the same spread. The wider range of centered feature $1$ remains visible. To check your axis choice, add the centered entries in each column: each sum should be zero.

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

Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for from numpy to jax.numpy.

import jax.numpy as jnp
x = jnp.array([[1., 10.], [3., 14.], [5., 18.]])

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

Step 2: Apply the core JAX transformation

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

mean = x.mean(axis=0)
centered = x - mean

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

Step 3: Verify shapes and numerical invariants

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

assert jnp.allclose(mean, jnp.array([3., 14.]))
assert jnp.allclose(centered.mean(axis=0), jnp.zeros(2))

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

Run the example

# From NumPy to jax.numpy: Feature preprocessing answers a specific question: how does each...
import jax.numpy as jnp
x = jnp.array([[1., 10.], [3., 14.], [5., 18.]])
# Reduce along axis=0 to compute `mean`.
mean = x.mean(axis=0)
centered = x - mean
print("Mean:", mean)
print("Centered:\n", centered)
assert jnp.allclose(mean, jnp.array([3., 14.]))
assert jnp.allclose(centered.mean(axis=0), jnp.zeros(2))

Mean: [ 3. 14.]
Centered:
 [[-2. -4.]
 [ 0.  0.]
 [ 2.  4.]]
new batch transformed: [[2.4494896 2.4494896]]
PASS: arrays-01

Check the normalization arithmetic

Predict before running: Predict the feature standard deviations and centered rows before comparing with NumPy.

# Experiment — Check the normalization arithmetic: Matching dtype and statistic definition makes this a useful...
import numpy as np
reference_data = np.array([[1., 10.], [3., 14.], [5., 18.]], dtype=np.float32)
# Reduce along axis=0 to compute `reference_mean`.
reference_mean = reference_data.mean(axis=0)
# Reduce along axis=0 to compute `reference_scale`.
reference_scale = reference_data.std(axis=0)
assert jnp.allclose(mean, reference_mean)
assert jnp.allclose(x.std(axis=0), reference_scale)
assert jnp.allclose(x.std(axis=0), jnp.sqrt(jnp.array([8. / 3., 32. / 3.])))

Matching dtype and statistic definition makes this a useful independent implementation check. It does not test TPU placement or performance.

Observe a new batch in training coordinates

Predict before running: Will a new row $[7, 22]$ transform to zero just because it is the only row in its batch?

# Experiment — Observe a new batch in training coordinates: The new observation is measured relative to training data.
def fit_standardizer(training_data):
    # Reduce along axis=0 to compute `fitted_mean`.
    fitted_mean = training_data.mean(axis=0)
    # Reduce along axis=0 to compute `fitted_std`.
    fitted_std = training_data.std(axis=0)
    # Combine or mask array elements to form `fitted_scale`.
    fitted_scale = jnp.where(fitted_std == 0., 1., fitted_std)
    return fitted_mean, fitted_scale
def apply_standardizer(data, statistics):
    fitted_mean, fitted_scale = statistics
    return (data - fitted_mean) / fitted_scale
statistics = fit_standardizer(x)
new_batch = jnp.array([[7., 22.]])
new_values = apply_standardizer(new_batch, statistics)
print("new batch transformed:", new_values)
assert jnp.allclose(new_values, jnp.full((1, 2), jnp.sqrt(6.)), atol=1e-6)

The new observation is measured relative to training data. Re-fitting on the one-row batch would erase this difference and produce constant-feature statistics.

Your exercise

Foundation · Standardize the original data using its standard deviation plus epsilon $10^{-6}$. Predict the transformed rows and check feature means and scales. Explain why the scale is only approximately one with epsilon.

# Exercise solution: Foundation · Standardize the original data using its standard...
# Reduce along axis=0 to compute `standardized`.
standardized = centered / (x.std(axis=0) + 1e-6)
assert jnp.allclose(standardized.mean(axis=0), 0., atol=1e-6)
assert jnp.allclose(standardized.std(axis=0), 1., atol=1e-5)

Test the constant-feature policy

Practice

Fit a standardizer to $[[1, 5]$, $[3, 5]$, $[5, 5]]$. Check the second training column and transform the new row $[7, 6]$. Explain why its second output is $1$ rather than zero.

Hint

Fit uses scale $1$ for a constant training feature. Applying statistics does not re-fit them.

# Test the constant-feature policy (Practice): The training feature is constant but the new value differs...
constant_data = jnp.array([[1., 5.], [3., 5.], [5., 5.]])
constant_statistics = fit_standardizer(constant_data)
constant_transformed = apply_standardizer(constant_data, constant_statistics)
# Confirm that all computed values remain finite (no NaN or Inf).
assert jnp.all(jnp.isfinite(constant_transformed))
assert jnp.allclose(constant_transformed[:, 1], 0.)
assert jnp.allclose(apply_standardizer(jnp.array([[7., 6.]]), constant_statistics)[0, 1], 1.)

Detect accidental batch-dependent preprocessing

Challenge

Transform two new rows together and separately using training statistics. The outputs should agree. Then show that re-fitting on the new batch changes the outputs.

Hint

The application function should depend on each row and fixed statistics, not on which other rows share the request.

# Detect accidental batch-dependent preprocessing (Challenge): This consistency test expresses an inference boundary: a...
new_rows = jnp.array([[7., 22.], [9., 26.]])
together = apply_standardizer(new_rows, statistics)
# Combine or mask array elements to form `separately`.
separately = jnp.concatenate([apply_standardizer(new_rows[i:i+1], statistics) for i in range(2)])
assert jnp.allclose(together, separately)
refitted = apply_standardizer(new_rows, fit_standardizer(new_rows))
assert not jnp.allclose(together, refitted)

Checkpoint

Why does x.mean(axis=0) return two values?

1. It averages each row
2. It preserves the two feature columns while reducing observations
3. JAX always returns two values

Diagnose the result

If a statistic has the wrong length, inspect the reduced axis. If normalized values become NaN, check zero scales and finite inputs before increasing epsilon. If a model behaves differently across inference batch sizes, inspect whether preprocessing re-fits statistics on each batch. If JAX and NumPy differ, align dtype, axis, and degrees-of-freedom choices before treating the difference as a library bug.

Carry forward

• Axes need semantic names in your reasoning; legal operations can still answer the wrong question.
• Fit preprocessing statistics on training data and reuse them for new observations.
• Constant and near-constant features need an explicit numerical policy.
• Independent arithmetic and matched NumPy checks establish the intended values.

Your evidence

Keep the fitted mean/scale, hand-derived training values, NumPy comparison, constant-feature test, and together-versus-separate new-batch results. Explain why re-fitting on each inference batch changes the problem.

Primary references

• JAX: arrays and NumPy interface
• JAX: standard deviation

Chapter 6 · arrays-02

Shapes, broadcasting, and dtypes

Specify broadcasting contracts and detect silent pairwise-loss and dtype mistakes.

What you will be able to do

• Apply the trailing-dimension broadcasting rule before execution.
• Distinguish feature-wise and observation-wise offsets with explicit singleton axes.
• Find a silent pairwise-residual bug using values and shape checks.
• State dtype and precision assumptions for numerical experiments.

The problem

Your code runs, but has it added the bias to the right axis? This is a common array puzzle: a legal operation can still answer the wrong question. We’ll use small tables to see exactly which values combine, then catch a broadcasting mistake that turns one error per example into a whole matrix of errors.

The idea

Before adding or subtracting arrays, name what each axis represents. Broadcasting aligns dimensions from the right and reuses singleton dimensions. A legal broadcast can still answer the wrong question, especially when a final mean hides the resulting shape.

Catch the scalar that hides a shape mistake

Take predictions $[1,3]$ and matching targets $[1,3]$. The intended residuals are $[0,0]$, so mean squared error is zero. If the targets are stored as a column, subtraction instead compares each target against both predictions. The residual matrix has entries $0,2,-2,0$, and its mean squared error is $2$.

Both calculations return a scalar after averaging. Looking only at the final loss, or checking that autodiff can differentiate it, will not expose the problem. Inspect the residual shape before reducing it and compare against a known aligned example.

This is also why a reshape should express a data contract. A singleton feature axis can mean one offset per observation; a singleton observation axis can mean one offset per feature. They are different operations even when an unlucky square fixture makes both expressions run.

Align axes before reducing

Predict: Why is checking that the loss is scalar insufficient?

Feature bias reuses a vector across observations; row offsets reuse each scalar across features. The final row creates all-pairs residuals. A mean would hide the unwanted $(n,n)$ intermediate. Each row describes a different shape contract.

Work the shape rule one dimension at a time

Align $(2, 3)$ and $(3,)$ as $(2, 3)$ and $(1, 3)$. The trailing $3$ matches, and the leading $1$ can broadcast to $2$. By contrast, $(2,)$ aligns as $(1, 2)$: trailing $2$ conflicts with $3$. That expression cannot add one scalar per row without another axis.

Represent row offsets as $(2, 1)$. The leading $2$ matches observations and the trailing $1$ broadcasts over the three features. A singleton axis is a statement about how data should be reused, not a random reshape to silence an exception.

feature bias: (2,3) + (1,3) → (2,3)
row offset:  (2,3) + (2,1) → (2,3)
wrong row offset: (2,3) + (2,) → incompatible

Do not test axes with only square arrays

If a batch happens to have $B=D$, a vector of length $B$ can accidentally fit as a feature vector. The expression succeeds while doing the wrong calculation. A test with two observations and three features forces those meanings apart.

Use distinguishable values, such as feature bias $[10, 20, 30]$ and row offsets $[100, 200]$. Then inspect at least two rows and columns. Shape checks alone do not detect every semantically wrong operation; the values need a known reference too.

A legal broadcast can change your loss

Predictions with shape $(B,)$ and targets with shape $(B, 1)$ broadcast to $(B, B)$. Each row then compares one target to every prediction. A mean over that matrix is a scalar, so a later gradient transform can accept it even though the objective is wrong.

For predictions $[1, 3, 5]$ and matching targets, the intended residual vector is all zero. A column target produces rows $[0, 2, 4]$, $[-2, 0, 2]$, $[-4, -2, 0]$, and a positive mean-squared error. Independent expected values expose the bug immediately. Validate prediction/target shape equality at the boundary where aligned examples are required.

(B,) − (B,1) → (B,B)
Intended aligned-example contract: prediction.shape == target.shape

Shape and dtype are separate contracts

Shape says how values are arranged; dtype says how each value is represented. Use explicit float32 inputs for this course experiment so the precision assumption is visible. Default JAX precision and X64 configuration can differ from expectations brought from a NumPy workflow.

A correctly shaped float32 calculation can still lose a small change at large magnitude. That is numerical representation, not an axis error. Record both shape and dtype in debugging output, and choose comparisons appropriate to their scale. Enabling a different dtype is not a substitute for correcting a broadcast.

Make boundary checks before reducing

A reduction often hides shape information. Before averaging residuals, check whether the array is the expected $(B,)$ vector or an unintended matrix. Likewise, check that feature statistics have one entry per feature and that row offsets have the observation dimension.

These checks belong at meaningful boundaries, rather than after every elementary addition. A descriptive failure before a loss is calculated is more useful than a plausible scalar whose semantics have already been lost.

Broadcasting repeats the bias across rows

Predict before looking: Why are both rows of the difference equal to $(10,20,30)$?

Read the figure

This heatmap shows the added amount, computed as the output minus the input batch. It does not show the final output array. Rows are observations and columns are features; the cell values and color bar measure added bias.

Both rows read $(10,20,30)$. The repeated vertical bands show that feature $0$ receives $10$ in every observation, feature $1$ receives $20$, and feature $2$ receives $30$.

Connect it to the computation

The bias has shape $(3,)$, which aligns with the last axis of the $(2,3)$ batch. Broadcasting applies the same feature-wise addition to both rows. For example, the input $(1,2,3)$ becomes $(11,22,33)$, while $(4,5,6)$ becomes $(14,25,36)$.

Identical rows in this picture mean identical additions, not identical predictions. If you intended a different offset for each observation, you would need a row-wise bias with a compatible shape, such as $(2,1)$; this figure would then have horizontal bands.

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

Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for shapes, broadcasting, and dtypes.

import jax.numpy as jnp
batch = jnp.array([[1., 2., 3.], [4., 5., 6.]], dtype=jnp.float32)

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

Step 2: Apply the core JAX transformation

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

bias = jnp.array([10., 20., 30.], dtype=jnp.float32)
y = batch + bias

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

Step 3: Verify shapes and numerical invariants

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

assert y.shape == (2, 3)
assert jnp.allclose(y[1], jnp.array([14., 25., 36.]))

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

Run the example

# Shapes, broadcasting, and dtypes: Before adding or subtracting arrays, name what each axis represents.
import jax.numpy as jnp
batch = jnp.array([[1., 2., 3.], [4., 5., 6.]], dtype=jnp.float32)
bias = jnp.array([10., 20., 30.], dtype=jnp.float32)
y = batch + bias
print(y)
print("Shape:", y.shape, "dtype:", y.dtype)
# Check tensor shape invariant: `y.shape == (2, 3)`
assert y.shape == (2, 3)
assert jnp.allclose(y[1], jnp.array([14., 25., 36.]))

[[11. 22. 33.]
 [14. 25. 36.]]
Shape: (2, 3) dtype: float32
aligned residuals: [0. 0. 0.]
pairwise residuals: [[ 0.  2.  4.]
 [-2.  0.  2.]
 [-4. -2.  0.]]
float32 large + 1 − large: 0.0
PASS: arrays-02

Make the pairwise bug visible

Predict before running: Matching predictions and targets should have zero loss. Predict the matrix created by a column target and its mean-square value.

# Experiment — Make the pairwise bug visible: The final result is scalar in both cases.
predictions = jnp.array([1., 3., 5.])
targets = jnp.array([1., 3., 5.])
correct_residuals = predictions - targets
pairwise_residuals = predictions - targets[:, None]
print("aligned residuals:", correct_residuals)
print("pairwise residuals:", pairwise_residuals)
# Check tensor shape invariant: `correct_residuals.shape == (3,)`
assert correct_residuals.shape == (3,)
# Check tensor shape invariant: `pairwise_residuals.shape == (3, 3)`
assert pairwise_residuals.shape == (3, 3)
assert jnp.allclose(jnp.mean(correct_residuals ** 2), 0.)
assert jnp.allclose(jnp.mean(pairwise_residuals ** 2), 16. / 3.)

The final result is scalar in both cases. Checking only the final loss shape would miss this objective bug.

Observe a precision limit separately

Predict before running: Does float32 preserve adding $1$ to $100$,$000$,$000$? Is the shape relevant to that result?

# Experiment — Observe a precision limit separately: This scalar example isolates precision from broadcasting.
large = jnp.array(100_000_000., dtype=jnp.float32)
print("float32 large + 1 − large:", float((large + 1.) - large))
assert float((large + 1.) - large) == 0.
# Check tensor shape invariant: `large.shape == ()`
assert large.shape == ()

This scalar example isolates precision from broadcasting. Autodiff and compilation do not make every arithmetic operation exact.

Your exercise

Add a different scalar offset to each row using offsets $[100., 200.]$. Make its shape explicit.

# Exercise solution: Add a different scalar offset to each row using offsets [100., 200.].
offsets = jnp.array([100., 200.])[:, None]
z = batch + offsets
# Check tensor shape invariant: `offsets.shape == (2, 1)`
assert offsets.shape == (2, 1)
assert jnp.allclose(z, jnp.array([[101.,102.,103.],[204.,205.,206.]]))

Use both kinds of offsets together

Practice

Add the feature bias and row offsets to the original batch. Derive every row before running, then check with an explicitly constructed array.

Hint

Use bias shape $(3,)$ and offsets shape $(2, 1)$. Do not swap their semantic roles.

# Use both kinds of offsets together (Practice): The first row receives 100 in every feature and the second...
combined = batch + bias + jnp.array([100., 200.])[:, None]
# Check tensor shape invariant: `combined.shape == (2, 3)`
assert combined.shape == (2, 3)
assert jnp.allclose(combined, jnp.array([[111., 122., 133.], [214., 225., 236.]]))

Reject an aligned-loss contract violation

Challenge

Write an MSE function that requires prediction and target shapes to match. Test zero loss for matching vectors and rejection of a column target. Explain why an automatic reshape might hide a data-ordering problem.

Hint

Raise a descriptive ValueError before subtracting. The check is based on shape metadata, not array truth values.

# Reject an aligned-loss contract violation (Challenge): Fix shapes at the data/model interface according to the...
def aligned_mse(prediction, target):
    # Guard input contract (`prediction.shape != target.shape`) and fail fast if violated.
    if prediction.shape != target.shape:
        raise ValueError("aligned prediction and target shapes must match")
    return jnp.mean((prediction - target) ** 2)
assert jnp.allclose(aligned_mse(predictions, targets), 0.)
# Run the boundary check and catch the expected exception:
try:
    aligned_mse(predictions, targets[:, None])
except ValueError:
    pass
else:
    raise AssertionError("Column targets must be rejected")

Checkpoint

Which offset shape adds one scalar to each row of a $(2, 3)$ batch?

1. $(2,)$
2. $(2, 1)$
3. $(3, 2)$

Diagnose the result

When broadcasting fails, align trailing dimensions on paper. When it succeeds with a surprising loss, inspect the residual shape before reduction and compare with hand-derived values. When a small change disappears, inspect dtype and magnitude. Use asymmetric shape tests so accidentally equal dimension sizes do not conceal a wrong axis.

Carry forward

• Broadcasting compatibility does not imply semantic correctness.
• A singleton axis declares reuse along a particular dimension.
• Prediction/target mismatch can silently create all-pairs residuals.
• Check dtype separately from shape and record numerical scale.

Your evidence

Keep an axis sketch, asymmetric offset checks, aligned-versus-pairwise residual arrays, a rejected shape mismatch, and the float32 precision observation. Explain the intended residual contract.

Primary references

• NumPy: broadcasting rules
• JAX: default dtypes and X64

Chapter 7 · arrays-03

Pure functions and explicit inputs

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

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.

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.

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

Step 1: Set up imports and input tensors

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

import jax.numpy as jnp

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

Step 2: Apply the core JAX transformation

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

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

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

Step 3: Verify shapes and numerical invariants

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

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

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

Run the 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.]))

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

Expose a mutable container alias

Predict before running: If alias = params, does assigning a new weight through alias preserve the original dictionary? Predict before running on a separate example.

# 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"]))

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

Replay two independent model configurations

Predict before running: What should each parameter setting predict, regardless of which was called first?

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

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

Your exercise

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

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

Hint

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

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

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.

Hint

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

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

Checkpoint

What makes predict easier to transform and test?

1. Reading an editable global weight
2. Passing parameters and data explicitly
3. Printing inside every numerical operation

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.

Your evidence

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.

Primary references

• JAX: pure functions and common gotchas
• JAX: stateful computations

Chapter 8 · arrays-04

Immutable updates and indexing

Implement functional indexed updates and distinguish accumulated contributions from shape-preserving masks.

What you will be able to do

• Predict the values returned by functional indexed set and add.
• Distinguish array-value semantics from Python name rebinding and compiler storage choices.
• Accumulate repeated-index contributions with an explicit integer reference.
• Choose a fixed-shape mask when output shape must remain stable.

The problem

You want to replace a few array entries, but JAX rejects the NumPy-style assignment. The key idea is small: describe the new array value, then keep the result. We’ll try replacement and addition, check that the original value is still available, and see what happens when an index appears more than once.

The idea

An indexed update in JAX returns an array value. This makes the relationship between old and new state explicit: the caller chooses which result to keep. The programming model describes values; compiler decisions about physical memory reuse are a separate issue.

Separate a new value from a new physical buffer

Start with $[2,4,6]$. Setting the middle entry to $10$ produces $[2,10,6]$. Adding $10$ at that position produces $[2,14,6]$. Both operations begin from the same original values; the second is not automatically applied to the result of the first.

Draw two arrows from the original array, one labeled set and one labeled add. If you want sequential updates, explicitly feed the first result into the second operation. This is the same state-flow idea we will use in optimization and recurrent loops.

Do not infer allocation cost from that drawing. A compiler may reuse storage when it can preserve the program's semantics. First verify old/new values and duplicate-index behavior for the operation you use, then measure memory behavior separately.

A Python name and an array value are different

After y = x.at[1].set(20), $x$ still denotes the original array and $y$ denotes the updated value. After x = x.at[1].set(20), the Python name $x$ is rebound to the new value. That spelling does not retroactively change another name holding the original value.

Keep original = x before an update when you want a reference for tests. Comparing original and updated makes the semantics visible. Repeatedly inspecting only the newly rebound name can make a functional update look like mutation even though the old value still exists.

Set replaces; add accumulates

Replacing an entry means its old value no longer contributes at that position in the result. Adding a contribution means the old value remains and the contribution is combined with it. Use set for a known replacement and add for a scatter-like accumulation.

Repeated indices need particular care. With add, contributions to the same index are accumulated; their floating-point order may not be guaranteed. With conflicting set values for the same index, do not assume a portable last-write-wins order. Choose unique indices when you need deterministic replacement semantics and test duplicate accumulation with exact small integers before studying floating-point ordering.

Choose a mask when you want to keep the shape

Selecting $x$[$x$ > threshold] extracts a variable number of elements: its output length depends on values. A masked replacement, jnp.where(x > threshold, x, 0), keeps the original shape. They are different calculations even if both are described casually as filtering.

A fixed-shape representation is useful for staged computation. Later lessons will explain why value-dependent output lengths can fail under jit. Decide whether downstream code expects compacted selected data or an array with masked entries before choosing the operation.

Make the index domain part of the contract

Our examples use explicit in-bounds integer indices and ordinary slices. Indexing behavior at out-of-bounds positions is not a substitute for validating bad input. Negative indexing, duplicate indices, masks, and scatter modes all need deliberate interpretation in a real data pipeline.

A small update test should verify unchanged positions as well as changed positions. Testing only the new value at one index could miss an unintended broadcast across a slice or a mistaken axis in a matrix.

Carry functional updates into training state

A model update is a larger version of the same idea: parameters and gradients produce new parameters. The caller retains or checkpoints state. Functional indexed operations let you describe a selective change without mutating an old array value.

This lesson prepares you to reason about optimizer and simulation state. It is separate from memory efficiency or a training algorithm; those require broader computations and measurements.

An update returns a new array

Predict before looking: Which cells change after set, and which after add?

Read the figure

Columns are array indices. The top row is the original array, and each lower row is a separate result created from that original. The cell colors encode stored values; the dark cell containing $20$ makes the replacement easy to locate.

The middle row replaces only index $1$: $(1,2,3,4)$ becomes $(1,20,3,4)$. The bottom row instead adds $5$ at indices $0$ and $2$, giving $(6,2,8,4)$.

Connect it to the computation

The bottom row still contains $2$ at index $1$. This is the visual clue that the two updates branch from the original; the second operation did not continue from the middle row. In both cases, the top row remains unchanged.

An immutable update returns a new value. To make changes accumulate, pass the returned array into the next update. If these operations were chained, the final array would be $(6,20,8,4)$, which is deliberately not the bottom row shown here.

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

Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for immutable updates and indexing.

import jax.numpy as jnp
x = jnp.array([1., 2., 3., 4.])

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

Step 2: Apply the core JAX transformation

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

y = x.at[1].set(20.)
z = x.at[jnp.array([0, 2])].add(5.)

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

Step 3: Verify shapes and numerical invariants

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

assert jnp.allclose(x, jnp.array([1.,2.,3.,4.]))
assert jnp.allclose(y, jnp.array([1.,20.,3.,4.]))
assert jnp.allclose(z, jnp.array([6.,2.,8.,4.]))

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

Run the example

# Immutable updates and indexing: An indexed update in JAX returns an array value.
import jax.numpy as jnp
x = jnp.array([1., 2., 3., 4.])
y = x.at[1].set(20.)
z = x.at[jnp.array([0, 2])].add(5.)
print("Original:", x)
print("Set:", y)
print("Add:", z)
assert jnp.allclose(x, jnp.array([1.,2.,3.,4.]))
assert jnp.allclose(y, jnp.array([1.,20.,3.,4.]))
assert jnp.allclose(z, jnp.array([6.,2.,8.,4.]))

Original: [1. 2. 3. 4.]
Set: [ 1. 20.  3.  4.]
Add: [6. 2. 8. 4.]
PASS: arrays-04

See name rebinding without old-value mutation

Predict before running: If two names point at the original array and one name is rebound, which values remain available?

# Experiment — See name rebinding without old-value mutation: Assigning a result changes the binding, while ignoring a result...
original = x
rebound = x
rebound = rebound.at[1].set(20.)
assert jnp.allclose(original, jnp.array([1., 2., 3., 4.]))
assert jnp.allclose(rebound, jnp.array([1., 20., 3., 4.]))
ignored = x.at[0].set(99.)
assert jnp.allclose(x, original)
assert float(ignored[0]) == 99.

Assigning a result changes the binding, while ignoring a result leaves the caller using the old value.

Accumulate repeated integer contributions

Predict before running: Predict the total at index $1$ after adding contributions $2$ and $3$ there. What differs from replacement?

# Experiment — Accumulate repeated integer contributions: The integer loop reference makes accumulation explicit.
import numpy as np
indices = jnp.array([1, 1, 3])
contributions = jnp.array([2, 3, 7], dtype=jnp.int32)
base_counts = jnp.zeros(4, dtype=jnp.int32)
counts = base_counts.at[indices].add(contributions)
# Allocate initialized array `reference_counts` with the specified shape and dtype.
reference_counts = np.zeros(4, dtype=np.int32)
for index, contribution in zip([1, 1, 3], [2, 3, 7]):
    # Accumulate the next contribution into `reference_counts[index]`.
    reference_counts[index] += contribution
assert jnp.array_equal(counts, reference_counts)
assert jnp.array_equal(counts, jnp.array([0, 5, 0, 7]))
assert jnp.array_equal(base_counts, jnp.zeros(4, dtype=jnp.int32))

The integer loop reference makes accumulation explicit. It avoids asserting a floating-point reduction order or conflicting-set behavior.

Your exercise

Replace the last two elements with zero in a new array. Keep $x$ unchanged.

# Exercise solution: Replace the last two elements with zero in a new array.
masked = x.at[2:].set(0.)
assert jnp.allclose(masked, jnp.array([1.,2.,0.,0.]))
assert jnp.allclose(x, jnp.array([1.,2.,3.,4.]))

Replace a region of a matrix

Practice

Create a $(3, 3)$ matrix containing $0$ through $8$. Set its final column to $-1$ in a returned value. Verify the whole result and original matrix.

Hint

The slice [:, $-1$] selects every row and the last column. A scalar replacement broadcasts across that slice.

# Replace a region of a matrix (Practice): Checking the entire matrix verifies both changed and...
# Construct and reshape `matrix` into the target tensor dimensions.
matrix = jnp.arange(9, dtype=jnp.float32).reshape(3, 3)
updated_matrix = matrix.at[:, -1].set(-1.)
assert jnp.allclose(updated_matrix, jnp.array([[0., 1., -1.], [3., 4., -1.], [6., 7., -1.]]))
assert jnp.allclose(matrix, jnp.arange(9).reshape(3, 3))

Keep shape while masking values

Challenge

Keep only values greater than $2$ from $x$, replacing other positions with zero. Compare with compact selection and explain the shape difference.

Hint

Use where for the fixed-shape result. Boolean indexing produces selected values, not a shape-preserving mask.

# Keep shape while masking values (Challenge): Both are useful representations, but they are not...
fixed_shape = jnp.where(x > 2., x, 0.)
selected = x[x > 2.]
# Check tensor shape invariant: `fixed_shape.shape == (4,)`
assert fixed_shape.shape == (4,)
# Check tensor shape invariant: `selected.shape == (2,)`
assert selected.shape == (2,)
assert jnp.allclose(fixed_shape, jnp.array([0., 0., 3., 4.]))
assert jnp.allclose(selected, jnp.array([3., 4.]))

Checkpoint

After y = x.at[1].set(20.), what happens to $x$?

1. Its second element becomes $20$
2. Its value remains unchanged
3. It becomes a Python list

Diagnose the result

If nothing changes, check whether the returned update was discarded. If old values seem changed, distinguish name rebinding from mutation and inspect mutable containers around the array. If duplicate replacements are surprising, do not assume last-write-wins order; use an explicit unique-index policy. If a mask later fails under jit, inspect whether its output length depends on runtime values.

Carry forward

• Indexed updates return values; retaining them is the caller’s responsibility.
• Name rebinding and mutation are different operations.
• Repeated-index add accumulates contributions, while conflicting set needs a deliberate policy.
• A shape-preserving mask and compact selection express different computations.

Your evidence

Keep original and updated arrays, a repeated-index integer reference, the whole-matrix replacement check, and compact-versus-fixed-shape outputs. Explain replacement, accumulation, and name rebinding.

Primary references

• JAX: indexed update API
• JAX: immutability and array updates

Chapter 9 · first-gradient

Your first gradient

Explain local sensitivity and the chain rule, then compare analytic, numerical, and automatic derivatives.

What you will be able to do

• Predict the sign and magnitude of a derivative from a local change.
• Trace the chain rule through a composed numerical function.
• Compare an analytic derivative, a finite-difference estimate, and JAX autodiff.
• Distinguish a stationary point from a minimum and repair scalar-output/input-type errors.

The problem

Imagine that you can turn one dial, and a score tells you how far you are from the result you want. Which way should you turn it? We’ll begin with the curve $f(x)=x^2$, where we can check the answer by hand. A derivative tells us how the score changes nearby; then JAX will calculate that derivative for us. You do not need to guess what the tool is doing.

The idea

A derivative tells us how an output changes near a particular input. We can understand it before using autodiff: move the input a small amount, estimate the output change from the local slope, then compare that prediction with the function.

Use a slope to predict a nearby change

For a separate hand calculation, let $f(x)=x^2$. At $x=3$, the derivative is $6$. Moving to $3.01$ gives a first-order prediction of $9+6(0.01)=9.06$. The exact value is $9.0601$; the small difference comes from curvature.

In the tangent figure, the tangent agrees with the curve at the selected point and has the same local slope. It need not follow the curve far away. Move the point and watch the slope change; then move only a small distance from that point to judge the linear approximation.

Autodiff computes the derivative of the operations you wrote. It does not choose a meaningful objective for you. A gradient of the wrong loss can be numerically correct, which is why later lessons check both the formula and its derivative.

Read a slope before you compute it

Let’s start at $x=3$. If we move right by $\delta=0.01$, the score changes from $9$ to $9.0601$. That is an increase of $0.0601$. The slope predicts an increase of about $6\times0.01=0.06$, which is close.

Here $\delta$ means a small change in the input, and $f\prime(x)$ means the slope at that input. Expanding the square gives the exact change below. The term $\delta^2$ is what the straight-line approximation leaves out. As the change gets smaller, that term becomes small compared with the linear term when the slope is nonzero.

Before running the code, try a negative input. At $x=-3$, moving a little to the right reduces the score. The slope should therefore be negative. At zero, the slope is zero even though moving away in either direction increases the square.

Three ways to obtain a derivative

An analytic derivative comes from reasoning about the mathematical function. A finite difference approximates a derivative by evaluating the function at nearby inputs. Automatic differentiation applies derivative rules to the operations that execute in the program and composes those rules with the chain rule.

Autodiff is not estimating a slope with a hidden step size, and it is not a symbolic algebra system simplifying a formula. Your program still runs floating-point operations. Those operations can overflow, lose precision, or evaluate a function whose mathematical derivative is not well-defined at the chosen point. Agreement among independent approaches is stronger evidence than merely printing one number.

Follow the chain rule through the program

Now split the calculation into two small steps. First compute $u=3x+1$, then square it to get $y=u^2$. A small change in $x$ first changes $u$; that change in $u$ then changes $y$. The chain rule multiplies the two local slopes.

The first slope is $3$. The second is $2u$. Multiplying gives $6(3x+1)$. At $x=2$, we have $u=7$, so the derivative is $42$. This gives us a hand-calculated answer for the JAX check. Autodiff follows the operations in the program; it is not estimating the derivative by repeatedly nudging the input.

Know the contract of grad

The default grad transformation differentiates the first positional argument. Its differentiated inputs must have an appropriate inexact dtype, such as a floating-point array, and its output must be scalar: shape $()$, not shape $(1,)$. A Python float like $3.0$ works; the integer $3$ is a different input type.

A vector of residuals is useful model output, but it does not yet specify a single objective. Reducing the residuals to a sum or a mean chooses an objective and its scaling. A Jacobian is the more appropriate object when you really want the derivatives of every vector output. We will first learn scalar objectives, then return to batching and derivative structure.

A zero gradient is a question, not a verdict

Both $x^2$ and $x^3$ have derivative zero at $x=0$. The square has a minimum there; the cubic keeps decreasing toward the left. The constant function $x\mapsto5$ also has zero derivative everywhere. A gradient is local first-order information, so it does not by itself prove optimality, global quality, or a correct implementation.

When a model produces a surprising zero gradient, ask whether the loss actually depends on the parameter, whether you differentiated the intended argument, and whether the chosen point is stationary. Checking nearby points and a simple analytic case is often more informative than immediately changing the optimizer.

A tangent describes local change

Predict before looking: At $x=3$, is the slope equal to the height of the curve?

Read the figure

The solid curve is $f(x)=x^2$. The dashed line is its tangent at the marked point $(3,9)$. The horizontal axis measures the input; the vertical axis shows either the function value or the tangent’s predicted value.

Near the marked point the two almost coincide. Farther away they separate: at $x=2$, the curve is $4$ and the tangent is $3$. At $x=0$, the curve is $0$, but the tangent has fallen to $-9$.

Connect it to the computation

The derivative at $x=3$ is $6$, so the local prediction is $f(3+\Delta x)\approx9+6\Delta x$. The exact value also contains $(\Delta x)^2$. That omitted term explains why the approximation is good for small changes and increasingly poor for large ones.

The negative part of the dashed line does not mean a square became negative. It means a local linear approximation was extended beyond the neighborhood where it is useful. A gradient describes local change, not the complete function.

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

Step 1: Set up imports and input tensors

Import the required JAX modules and define the initial inputs for your first gradient.

import jax
import jax.numpy as jnp

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

Step 2: Apply the core JAX transformation

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

def f(x):
    return x ** 2
# Differentiate the objective to obtain `derivative` via automatic differentiation.
derivative = jax.grad(f)

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

Step 3: Verify shapes and numerical invariants

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

assert jnp.allclose(derivative(3.0), 6.)
assert jnp.allclose(derivative(0.0), 0.)

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

Build a numerical estimate

# Numerical estimate: Approximately 6.0.
def central_difference(function, x, h=1e-3):
    return (function(x + h) - function(x - h)) / (2 * h)

def square(x):
    return x * x

estimate = central_difference(square, 3.0)
print("Finite-difference estimate:", estimate)
assert abs(estimate - 6.0) < 1e-8

Approximately $6.0$. This estimate evaluates the function at two nearby points; it is not automatic differentiation.

Run the example

# Your first gradient: A derivative tells us how an output changes near a particular input.
import jax
import jax.numpy as jnp
def f(x):
    return x ** 2
# Differentiate the objective to obtain `derivative` via automatic differentiation.
derivative = jax.grad(f)
print(float(derivative(3.0)))
assert jnp.allclose(derivative(3.0), 6.)
assert jnp.allclose(derivative(0.0), 0.)

Finite-difference estimate: 5.999999999999339
6.0
chain rule: -1.0 -12.0
chain rule: 0.0 6.0
chain rule: 2.0 42.0
h / estimate / error: 0.1 27.009973526000977 0.009973526000976562
h / estimate / error: 0.01 27.000045776367188 4.57763671875e-05
h / estimate / error: 0.001 26.9975643157959 0.0024356842041015625
h / estimate / error: 1e-05 27.0843505859375 0.0843505859375
h / estimate / error: 1e-07 0.0 27.0
PASS: first-gradient

Check the chain rule at three inputs

Predict before running: For $g(x)=(3x+1)^2$, predict the derivative at $-1$, $0$, and $2$. Explain why the sign changes.

# Experiment — Check the chain rule at three inputs: The inner derivative contributes a factor of 3.
def composed(x):
    return (3. * x + 1.) ** 2
for point in (-1., 0., 2.):
    # Differentiate the objective to obtain `observed` via automatic differentiation.
    observed = jax.grad(composed)(point)
    expected = 6. * (3. * point + 1.)
    print("chain rule:", point, float(observed))
    assert jnp.allclose(observed, expected)

The inner derivative contributes a factor of $3$. Omitting that factor is a chain-rule error; matching only one point can hide a wrong expression.

Choose a finite-difference scale

Predict before running: For the float32 function $x^3$ at $x=3$, will shrinking $h$ forever keep improving the estimate? Run the sweep and compare with $27$.