I Built a Neural Network Training Loop in 5 Lines Using JAX - Day 2

SERIES: Learning RL and JAX in Public - from zero to DeepMind :)

On Day 1 I covered what JAX is and the three superpowers: GPU acceleration, grad, and jit. Today I actually used them. And the moment where I manually verified a gradient by hand and watched JAX confirm it is the one that made everything real for me.

But before the code, I had to understand what a gradient actually is. Because I realised I had been using them for years without a clean mental model.

What is a derivative, really?

Imagine you are adjusting a dial on a machine. The dial controls the volume. Your goal is to get the volume to exactly 50. Right now it is at 30. You turn the dial a tiny bit. The volume goes up. The derivative answers one question: if I nudge this dial by a tiny amount, how much does the output change? That is it.

In math:

• f(x) = x * x
• f(3) = 9
• f(3.001) = 9.006001
• nudge of 0.001 caused a change of 0.006
• rate of change = 0.006 / 0.001 = 6

The derivative of x squared at x equals 3 is 6. Meaning: at this point, if x increases by a tiny amount, the output increases 6 times as fast.

What is a gradient?

A derivative is for a function with one input. A gradient is the same idea for a function with many inputs. It is just a list of derivatives, one per input.

• function with 1 input → derivative (a single number)
• function with N inputs → gradient (a list of N numbers)

Why does this matter for ML?

Every neural network has weights. Training means finding the weights that make the model least wrong. "How wrong" is measured by a number called the loss. The gradient of the loss tells you: for each weight, does increasing it make the loss bigger or smaller, and by how much?

Nudge every weight in the direction that reduces the loss. Repeat thousands of times. That is gradient descent. That is literally how every neural network on the planet trains.

Now the code

First I verified grad by hand:

import jax

def square(x):
    return x * x

grad_of_square = jax.grad(square)
print(grad_of_square(3.0))  # 6.0

Derivative of x squared is 2x. At x equals 3, that is 6. JAX says 6.0. Verified.

Then I wrote a real loss function:

def loss(w):
    x = 2.0
    target = 10.0
    prediction = w * x
    return (prediction - target) ** 2

grad_of_loss = jax.grad(loss)

This loss measures how wrong a prediction is. The ideal weight is 5.0 because 5 times 2 equals 10.

Then gradient descent from scratch:

learning_rate = 0.1
w = 1.0

for step in range(5):
    grad = grad_of_loss(w)
    w = w - learning_rate * grad
    print(f"step {step + 1}: w = {w:.4f}, loss = {loss(w):.4f}")

Output:

step 1: w = 2.6000, loss = 21.1600
step 2: w = 3.6640, loss = 8.6491
step 3: w = 4.3302, loss = 3.5369
step 4: w = 4.7481, loss = 1.4461
step 5: w = 5.0105, loss = 0.0044

(almost zero, w approaching 5.0)

That loop is the core of every neural network training run. More weights, more complex functions, but the same mechanic.

The jit speedup

I benchmarked a large computation with and without jit over 100 runs:

def slow_computation(x):
    return jnp.sum(jnp.sin(x) ** 2 + jnp.cos(x) ** 2)

fast_computation = jax.jit(slow_computation)

Result: 32x speedup on my MacBook CPU. On a GPU this goes to 100x or more.

The first call is always slower because that is when JAX compiles. Every call after uses the cached compiled version. One catch: if you change the shape of your input, JAX recompiles. So keep your array shapes consistent.

One question I had: do I need to memorise all this syntax?

No. The important thing is knowing what to use and when. Nobody in a research role writes jax.grad from memory every time - they know it exists, they know what it does, they look up the exact syntax when needed.

What you do need to carry in your head: grad differentiates, jit compiles, and they compose together. That mental model handles 90% of situations.

Day 3 tomorrow: vmap, the fourth JAX superpower. This one genuinely surprised me.

All code from this series, organised by day, is on my GitHub: https://github.com/MadhumithaKolkar/jax-rl-lab

Happy learning everyone!
Madhumitha Kolkar (index_0)