Before JAX existed, a small Python library called Autograd let researchers compute gradients of native Python and NumPy code automatically. It was elegant: write normal math, get derivatives for free. But Autograd ran on CPUs only, couldn't compile code for speed, and couldn't scale to modern accelerators like GPUs and TPUs.
In 2018, a team at Google Brain — including Matt Johnson, Roy Frostig, Dougal Maclaurin, and Chris Leary — set out to build Autograd's spiritual successor. Their goal was ambitious: take the simplicity of "differentiate Python functions" and combine it with Google's XLA compiler, which could target GPUs, TPUs, and future hardware. The result was JAX — originally standing for "Just After eXecution," though the name now simply stands on its own.
JAX was open-sourced in late 2018, and it took the research community by storm. Unlike TensorFlow's graph-building ceremony or PyTorch's object-oriented layers, JAX offered something different: composable function transformations. You write a plain Python function using NumPy-like operations, then transform it — compile it, differentiate it, vectorize it, parallelize it — all by wrapping it with simple function calls.
💡 Why This Matters
JAX wasn't built to be "another deep learning framework." It was built to be a numerical computing system that happens to be great for ML. This distinction explains why JAX feels different from PyTorch or TensorFlow — it's a tool for researchers who think in terms of mathematical functions, not neural network layers.
Today, JAX is the primary research framework at Google DeepMind. It powered the training of AlphaFold (the protein structure prediction system), research behind Gemini, and countless research papers. Major research labs worldwide — including teams at Stanford, MIT, Berkeley, and MILA — use JAX for work in scientific computing, physics simulations, Bayesian inference, and reinforcement learning.
Here's what a first encounter with JAX looks like:
import jax
import jax.numpy as jnp
# A plain function — no classes, no decorators needed
def f(x):
return jnp.sum(x ** 2)
x = jnp.array([1.0, 2.0, 3.0])
# Evaluate it
print(f(x)) # 14.0
# Now get its gradient — just wrap the function
grad_f = jax.grad(f)
print(grad_f(x)) # [2. 4. 6.]
# Compile it for speed
fast_f = jax.jit(f)
print(fast_f(x)) # 14.0 (but faster on subsequent calls)
That's JAX's core promise: you write functions, and JAX gives you tools to transform them. No framework classes to inherit from, no special tensor types to learn. Just functions in, functions out.