JAX: Compose Gradients, Compilation, and Vectorization Like Building Blocks

You've written a NumPy function that works fine on your laptop, but now you need gradients—and you need them fast, on a GPU or TPU. Rewriting everything in a lower-level framework feels like overkill. JAX is a Python library for accelerator-oriented array computation and program transformation, and it's built around the idea that you shouldn't have to choose between writing natural Python and getting high-performance numerical computing.

What It Does

JAX is a Python library for high-performance numerical computing and large-scale machine learning. At its core, it's an extensible system for transforming numerical functions—and the transformations are the whole point. The README highlights three: jax.grad for automatic differentiation, jax.jit for compilation, and jax.vmap for vectorization.

JAX can automatically differentiate native Python and NumPy functions. That includes differentiating through loops, branches, recursion, and closures—and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (backpropagation) via jax.grad, forward-mode differentiation, and you can compose the two arbitrarily to any order.

For performance, JAX uses XLA to compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators. You compile pure functions with jax.jit, and—this is the key part—compilation and automatic differentiation can be composed arbitrarily. Under the hood, JAX is really a system for composable function transformations at scale.

Why It's Cool

• The transformations compose, and that's the real feature. The README's opening example stacks three transformations in a single line: jax.jit(jax.grad(loss)) compiles a gradient function, and then jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0))) wraps that in a vectorization to get fast per-example gradients. You're not picking one trick—you're nesting them.

• Higher-order derivatives come for free. Because grad returns a function, you can just call grad on it again. The README shows jax.grad(jax.grad(jax.grad(tanh)))(1.0) printing a third derivative. No special API, no ceremony.

• Python control flow still works. A plain if x > 0 branch inside your function doesn't break differentiation—the README's abs_val example shows jax.grad handling it correctly at both positive and negative inputs. That's a meaningful convenience if you've fought with tracing restrictions elsewhere.

• It's honest about being a research project. The README says plainly: "This is a research project, not an official Google product. Expect sharp edges." That kind of framing tells you what you're getting into—and it points you to a dedicated gotchas document rather than pretending the rough parts don't exist.

How to Try It

JAX is on PyPI, so installation is a single command:

pip install jax

Then you can start with something small. Here's the README's tanh example, which shows differentiation and higher-order differentiation in a few lines:

import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0 * x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
# prints 0.4199743

print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
# prints 0.62162673

Once that clicks, try composing transformations on a real function. The README's loss-function example is a good next step—it builds a compiled gradient with jax.jit(jax.grad(loss)) and then vectorizes it with jax.vmap to get fast per-example gradients.

If you hit something confusing, the README links to the Common Gotchas notebook and the Autodiff Cookbook. The full reference docs live at docs.jax.dev, and the source is at github.com/jax-ml/jax.

Final Thoughts

JAX is best suited to people doing numerical computing or machine learning research who want fine-grained control over how their functions are differentiated, compiled, and vectorized—and who are comfortable with a tool that's still evolving. The composability is the draw: once you internalize that grad, jit, and vmap are just functions you can stack, a lot of previously awkward code gets shorter. It's not a polished product, and the README says so. But if you're willing to work with its rough edges, the transformation model is worth the trip.