DEV Community: Madhumitha Kolkar

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

Madhumitha Kolkar — Tue, 21 Jul 2026 04:24:25 +0000

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)

I Started Learning JAX as a Senior ML Engineer - Here's My First Impression.

Madhumitha Kolkar — Mon, 20 Jul 2026 05:41:20 +0000

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

I have been a machine learning engineer for six years. PyTorch is basically muscle memory at this point. So when I decided to seriously learn JAX, I expected it to feel like switching keyboards - same idea, different layout.

It is not that. JAX is a genuinely different way of thinking about numerical computing. And on Day 1, I want to share the mental model that made it click for me.

If you already know NumPy, you already know 80% of JAX's syntax. That part is intentional. The JAX team designed the API to mirror NumPy almost exactly so the learning curve stays shallow.

# NumPy
import numpy as np
x = np.array([1.0, 2.0, 3.0])
print(np.sum(x))  # 6.0

# JAX - identical
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
print(jnp.sum(x))  # 6.0

Same thing. Where JAX diverges is in three superpowers that NumPy simply does not have.

Superpower 1: It runs on GPU and TPU automatically

No .to('cuda'). No .cuda(). JAX looks for available accelerators and uses them. You write the same code whether you are on a laptop CPU or a cloud TPU. This alone is huge if you work across different environments.

Superpower 2: jax.grad - automatic differentiation as a function

In PyTorch, you call .backward() on a loss tensor. In JAX, grad is a standalone function transformer. You hand it any Python function, it hands you back a new function that computes the gradient of the original.

import jax

def square(x):
    return x ** 2

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

The derivative of x squared is 2x. At x equals 3, that is 6. JAX confirms it instantly. The fact that this works on any arbitrary Python function is what makes JAX powerful for research.

Superpower 3: jax.jit - compile your code for speed

Python is slow by default. jit (Just-In-Time compilation) takes your function, compiles it to optimised machine code on the first call, and every subsequent call skips Python entirely.

fast_fn = jax.jit(my_function)

First call is slower because of compilation. Every call after that is significantly faster. On Day 2 I benchmarked this and got a 32x speedup. More on that in the next post.

The one rule that surprised me: no in-place mutation

PyTorch lets you modify arrays in place. JAX does not. Every operation creates a new array. This feels wrong for about two days. Then you realise it is what makes grad and jit mathematically clean - they need functions to be pure (same input always gives same output). Mutation breaks that.

# This errors in JAX
x[0] = 99

# Do this instead
x = x.at[0].set(99)  # creates a new array

Why top research labs use JAX

The more I dig into JAX, the more I understand why the most ambitious ML research happening today is built on it. Labs like Google DeepMind run their entire research stack on JAX and Flax, and it is not arbitrary. When you are training models across hundreds of TPUs, you need code that is compiled, mathematically clean, and composable at scale. The combination of jit and grad doing exactly that is genuinely elegant. I find myself wanting to understand not just how to use these tools but why they were designed this way, and that curiosity is pushing me to go deeper into the fundamentals than I ever did with PyTorch. There is a certain level of ML engineering where understanding the infrastructure is just as important as understanding the models, and JAX sits right at that intersection.

Day 1 was arrays and basic operations. Simple stuff. But getting the mental model right before writing complex code matters more than rushing into the deep end.

Day 2 is grad, jit, and a gradient descent loop from scratch. See you there.

If you are on a similar journey from PyTorch to JAX, I would love to hear what tripped you up first. Drop it in the comments.

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)