JAX Pathways

Course · Setup & first stepsPhase 00 · Study guide

Can someone else reproduce your first result?

Onboarding · Local CPULaptop CPU + TPU bridge4 lessons (~3.2h)

Revision v0.7.0 ·

Start with one saved program and a known answer. The terminal, interpreter and working folder are part of the experiment: a notebook that imports a package is not yet proof that the downloaded script uses the same environment.

Preparation: Basic Python; no previous JAX experience.

Runtime: CPU; the TPU bridge is optional accelerator practice

Start the phaseSee the project

Tools, project folders and experiment organization: prepare a run directory and keep its configuration and evidence together.

Check your starting point

You saved a script in your course folder, but the terminal reports that the file does not exist. Which observation should you check first?

Compare your reasoning

Check the terminal’s current folder and the exact filename, including its extension. Then run the environment’s Python explicitly. Reinstalling JAX will not repair a wrong path.

If this is unfamiliar, work through Set up your learning workspace before attempting the project.

Build in stages

1. Make one reproducible result

Run the row-sum example as a saved file. Record the Python executable, package versions, backend, input shape and computed sums. Explain why a package version alone cannot identify which interpreter ran the file.

2. Separate local practice from parallel TPU execution

Start a fresh process for virtual CPU devices and compare its device receipt with the ordinary process. Then learn how to get a TPU running (hosted TPU or single-host Cloud TPU VM), launch your experiment over SSH with explicit backend checks, and run subsequent course phases on TPU in parallel.

Try a changed condition

Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.

A colleague gets the right sums but a different device count. Decide whether the arithmetic is wrong and how to compare the runs.

Compare an approach

Compare values and shapes first, then backend, process startup flags and environment. Device enumeration is a runtime property; equal arithmetic does not imply equal hardware or performance.

Diagnose a likely failure

Symptom: Import works in a notebook but fails in the terminal.

Check next: Compare interpreter paths before changing packages; launch both from the intended environment.

Decide what is ready

Keep a saved script, its command and actual output, plus an explanation of one repaired setup error. Use foundation-toolkit stage 1. A screenshot without a rerunnable file is incomplete.

Where to put your project code

Write your code in: projects/welcome-audit/starter/model.py (copy projects/welcome-audit/starter/model.py first: cp projects/welcome-audit/starter/model.py projects/welcome-audit/starter/model.py).

Verify your implementation: python3 projects/welcome-audit/tests/check.py --stage all --impl projects/welcome-audit/starter/model.py

Run the staged project checks on your implementation in projects/welcome-audit/starter/model.py and record your transfer experiment results in your portfolio.

Open Check the CPU environment and a saved matrix arithmetic identity

Open Connect arrays, transformations and explicit state

What still needs further work

Fresh-machine Windows/Linux/macOS walkthroughs and novice usability sessions remain unverified; accelerator setup needs a separately observed device run.