Onboarding · Local CPULaptop CPU + TPU bridgeLesson 2 of 4 in Phase 00

Phase 00 / Setup & first steps

Meet your arrays and devices

Revision v0.7.0 · 2026-10-20

Lesson · about 45 minutes · CPU practice

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

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Set up your learning workspace
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson text

Setup help: Python, notebook, or Cloud TPU VM

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

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

macOS / Linux terminal · repository folder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder

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

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

macOS / Linux terminal · activated environment

python phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py

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

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

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

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

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

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

Check: Shape: $(2, 3)$ ; Dtype: float32; Row sums: [ $3$ . $12$ .]. Device names vary.

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

Phase 00.02 Stage 1 of 5 · Understand

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

What you will be able to do

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.

Y_{i,j} = \alpha X_{i,j} + \beta, \qquad X \in \mathbb{R}^{M \times N}

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.

Open full-size figure Download data and provenance

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

Where to put & run this code: Work in phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py (or downloaded exercises/welcome-02.py / exercises/welcome-02.ipynb) and execute from the repository root with python3 phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py.

1 Prepare the inputs

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

main.py · append build step 1

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

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

2 Build the computation

Append this block below the inputs in the same file.

main.py · append build step 2

# Step 2 — Build the computation: 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)

Code building blocks in this step

jnp.arange(n, dtype=...) — Creates a 1-D JAX array of evenly spaced values [0, 1, ..., n-1] on the target device.
array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).

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

3 Run and check the result

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

main.py · append build step 3

# Step 3 — Run and check the result: Compare the output to the expected result below before making the...
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.]))

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

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

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

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

# Step 2 — Build the computation: 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.]))

Run the example

python3 phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py

About this command

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

Inspect the complete assembled example

main.py · complete Python example

# Step 1 — Prepare the inputs: These explicit inputs define the case that the later checks will...
import jax
import jax.numpy as jnp
# Step 2 — Build the computation: 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.]))

Expected result: Shape: $(2, 3)$ ; Dtype: float32; Row sums: [ $3$ . $12$ .]. Device names vary.

Recorded reference execution

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

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

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

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

Center each measurement

Center each measurement · Python experiment

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

Expected: Means $[[1.5,2.5,3.5]]$ ; centered rows are $-1.5$ and $+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

Keep a host reference · Python experiment

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

Expected: A NumPy array with the same two-by-three values.

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

Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py (or inside the Make it yours cell of exercises/welcome-02.ipynb) and run python3 phases/00-welcome/02-meet-your-arrays-and-devices/code/main.py to verify your prediction before opening the reference solution.

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

Show a writing guide and starter scaffold

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

1. Key functions & syntax to use:

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

2. Step-by-step implementation plan:

Reduce along axis=0 to compute column_sums.
Check tensor shape invariant: column_sums.shape == (3,)

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

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

Reveal the reference solution

Reference solution

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

Show a writing guide and starter scaffold

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

1. Key functions & syntax to use:

jnp.arange(n, dtype=...) — Creates a 1-D JAX array of evenly spaced values [0, 1, ..., n-1] on the target device.
jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).

2. Step-by-step implementation plan:

Construct and reshape table into the target tensor dimensions.

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

# 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(...)  # TODO: compute table
assert jnp.allclose(table.sum(axis=1),jnp.array([3.,12.,21.]))  # TODO: complete assertion check
assert jnp.allclose(table.mean(axis=0),jnp.array([3.,4.,5.]))  # TODO: complete assertion check

Show a hint

Keep observation rows and measurement columns.

Reveal reference solution and reasoning

Add a new observation · Python solution

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

Changing the observation count changes axis $0$ , while the measurement count stays three.

Diagnose an impossible reshape

Challenge

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

Show a writing guide and starter scaffold

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

1. Key functions & syntax to use:

array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).

2. Step-by-step implementation plan:

Run the boundary check and catch the expected exception:

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

# 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  # TODO: complete assertion check

Show a hint

Count required elements before changing the code.

Reveal reference solution and reasoning

Diagnose an impossible reshape · Python solution

# 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

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

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.

Keep your evidence

Download evidence template

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

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

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

Save work for this lesson

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

Meet your arrays and devices

What you will be able to do

The problem

The idea

Read one array two ways

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

Rows are groups of observations

Read the figure

Connect it to the computation

1 Prepare the inputs

2 Build the computation

3 Run and check the result

Run the example

Recorded reference execution

Center each measurement

Keep a host reference

Make it yours

Add a new observation

Diagnose an impossible reshape

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references