Part 8 of 12

JAX on the CPU: Field Notes with a Bot

Field notesBy Bot Tutorials · Published Nov 3, 2026 · Facts last verified Nov 3, 2026

ML libraries series complete. "Review Your Training Code with a Bot: Checklist" was the finale. See every part in the series hub.

Part 8: install JAX for CPU, take a gradient, compile with jit and use explicit random keys, with real output and a hand-checked answer.

Video: JAX on the CPU: Field Notes with a Bot (4:49). Narrated from this guide, with the real code and real output we ran. Captions included (caption file). Prefer reading? Everything in the video is in the guide below.

JAX takes NumPy-style code and adds two tricks: automatic gradients and compilation with jit. It's popular in research, and it's a good library for learning how gradients really work, because you call grad on an ordinary Python function. These field notes are what we saw installing and running JAX on a CPU-only box, with a bot explaining along the way.

What the bot can and can't do here

It can:

Explain what grad, jit and vmap do.
Convert a small NumPy function to jax.numpy.
Work through the calculus so you can check JAX's answer.

It can't:

Keep up with JAX's install extras. Check the installation page.
Know about JAX's sharp edges unless you ask, such as immutable arrays and explicit random keys.
Promise speed-ups from jit. Measure them.

Check first, compile second

Run a function normally, check its answer by hand, then add jit. If the answers differ, you've learned something before it cost you.

Install

JAX's installation page says it needs two packages: jax, which is pure Python, and jaxlib, which holds the compiled parts. For CPU, it gives:

pip install --upgrade pip
pip install --upgrade jax

On our box that installed jax and jaxlib 0.12.0.1. For NVIDIA GPUs the page recommends the pip extras, such as jax[cuda13], and says the driver must be 525 or later for CUDA 12 and 580 or later for CUDA 13 on Linux. On CUDA 13 it supports GPUs with SM version 7.5 or newer.

The script

import jax, jax.numpy as jnp
print("jax", jax.__version__, "devices", jax.devices())
def loss(w, x, y):
    return jnp.mean((w * x - y) ** 2)
x = jnp.arange(1.0, 6.0); y = 3.0 * x
g = jax.grad(loss)
print("grad at w=1:", float(g(1.0, x, y)))
fast = jax.jit(loss)
print("jit loss at w=3:", float(fast(3.0, x, y)))
key = jax.random.key(0)
print("random (key 0):", jax.random.normal(key, (3,)).round(3).tolist())

Real output:

jax 0.12.0.1 devices [CpuDevice(id=0)]
grad at w=1: -44.0
jit loss at w=3: 0.0
random (key 0): [1.6230000257492065, 2.0250000953674316, -0.43400001525878906]

Field note 1: check the gradient by hand

The loss is the mean of (w·x − y)² with x = 1, 2, 3, 4, 5 and y = 3x. Its derivative with respect to w is the mean of 2·x·(w·x − y). At w = 1, each term is 2·x·(x − 3x) = −4x². The squares add up to 55, so the mean is −4 × 55 ÷ 5 = −44. JAX printed -44.0. That's the habit: when a bot or a library gives you a number, check one by hand.

Field note 2: jit gives the same answer

jax.jit(loss) compiled the function, and at w = 3 the loss is 0, because 3x matches y exactly. Compiled and uncompiled versions should agree. We didn't time them here, so we make no claim about speed.

Field note 3: random numbers need keys

JAX doesn't use a hidden global random state. You make a key with jax.random.key(0) and pass it in. The same key gives the same numbers, so we saw 1.623, 2.025, -0.434 (rounded). To get new numbers you split the key. JAX's Quickstart and docs explain this. Bots trained on NumPy habits often forget it.

Field note 4: arrays are immutable

In NumPy you can write a[0] = 5. In JAX that raises an error, and you write a = a.at[0].set(5) instead. The jax.numpy reference lists this and other differences.

Prompts that worked

Convert this NumPy function to jax.numpy. List every change you made and why. If anything depends on in-place updates or global random state, rewrite it the JAX way.

Here's my loss function. Work out the derivative by hand at w = 1 for these inputs, then I'll compare it with jax.grad.

Words you'll see

grad: a function transformation that returns the gradient of a function.
jit: "just in time" compilation of a function for faster repeated calls.
vmap: automatic vectorization of a function over a batch dimension.
PRNG key: the explicit random state JAX passes around instead of a global seed.
Device: where arrays live. On our box, jax.devices() showed a single CpuDevice.

If a bot tells you JAX will be faster for your job, ask it how you would measure that on your own machine, and then measure it. A fair test runs each version a few times on the same data, and skips the first jit call, which includes compile time.

When to reach for JAX

JAX suits work where you want gradients of your own functions, or compiled numerical code that looks like NumPy. If you're following a course or a model that uses PyTorch or TensorFlow, stick with that library first. A bot can translate small examples between them, but check every translation by running both versions on the same input.

Tested vs documented

Tested by us: the CPU install, the gradient, jit result and random numbers, on JAX 0.12.0.1.
Documented, not tested by us: GPU installs and driver requirements, from JAX's installation page.
Not claimed: any speed-up from jit.

Where to go next

The last part today reads Hugging Face model cards and licenses before you download anything.

Safety, licensing and privacy: read this before you run anything

This series is about learning and small personal projects with machine learning libraries. A bot is a helper, not an authority.

Check bot advice against official docs. Version numbers, install commands and compatibility tables change often. Compare every answer with the PyTorch, TensorFlow, JAX, scikit-learn, NVIDIA, Hugging Face or ONNX Runtime pages we link. If the bot and the docs disagree, the docs win.
Work in a virtual environment. Don't install or remove packages in your system Python, and don't run sudo pip.
Run untrusted files carefully. Pickle-based model files can run code when loaded. Only load models from sources you trust, and prefer safer formats like safetensors where the library supports them.
Respect licenses. Datasets and model weights carry their own licenses and terms. Read the model card and license before you use, share or sell anything built on them. We don't give legal advice; ask a qualified person if it matters.
Protect personal data. Don't paste private, customer or medical data into a bot or a training set unless you have the right to, and check the bot's own privacy settings first.
No hardware claims you didn't measure. Treat any speed figure from a bot as a guess until you've timed it on your own machine.
No medical, legal or financial decisions from a toy model. These guides are for learning only.

FAQ

Is JAX a drop-in replacement for NumPy?

jax.numpy looks like NumPy, but JAX arrays are immutable and random numbers use explicit keys. Check the jax.numpy reference for differences.

How do I install JAX for a GPU?

JAX's installation page recommends pip wheels such as jax[cuda13], and says the driver must be 580 or later for CUDA 13 on Linux. We didn't test GPU installs.

Was this tested?

Yes, on the CPU. The gradient, jit result and random numbers are real output from JAX 0.12.0.1 on our box.