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

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

Enter fullscreen mode Exit fullscreen mode

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

Enter fullscreen mode Exit fullscreen mode

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)

Enter fullscreen mode Exit fullscreen mode

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

Enter fullscreen mode Exit fullscreen mode

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)