The Core Thesis: Composable Function Transformations

Most ML frameworks are built around objects: you create a model object, feed data into it, and call methods on it. JAX takes the opposite approach. Its central idea is composable function transformations — higher-order functions that take a function as input and return a new, transformed function as output.

Think of it like this: in mathematics, if you have a function f(x), you can create a new function f'(x) that gives you the derivative. You can also create a compiled version that runs faster, or a vectorized version that handles batches. JAX makes each of these a concrete operation you can call in Python.

The four foundational transformations — JAX's "four pillars" — are:

• jax.jit — Compile: traces your function and compiles it into optimized machine code via the XLA compiler. First call is slower (compilation), subsequent calls are dramatically faster.
• jax.grad — Differentiate: takes a scalar-valued function and returns a new function that computes its gradient. This is automatic differentiation in its purest form.
• jax.vmap — Vectorize: takes a function that works on a single example and returns a function that works on a batch. You never write batch dimension logic manually.
• jax.pmap — Parallelize: runs a function across multiple devices (GPUs/TPUs) simultaneously. Each device gets a slice of data and runs the same program.

The key word is composable. These transformations can be stacked freely:

import jax
import jax.numpy as jnp

def loss_fn(params, x, y):
    predictions = jnp.dot(x, params)
    return jnp.mean((predictions - y) ** 2)

# Compose transformations: differentiate, then compile
fast_grad = jax.jit(jax.grad(loss_fn))

# Or: vectorize per-example gradients, then compile
per_example_grads = jax.jit(jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0)))

params = jnp.array([1.0, 2.0])
x = jnp.array([[1.0, 0.5], [0.3, 0.8], [0.9, 0.1]])
y = jnp.array([1.5, 1.0, 0.8])

# Get compiled gradient
print(fast_grad(params, x, y))

# Get per-example gradients (one gradient vector per data point!)
print(per_example_grads(params, x, y))

Each transformation preserves the function's interface. jax.grad(f) returns a function with the same signature as f. jax.jit(f) returns a function with the same signature. This means you can always add another layer of transformation without refactoring your code.

This functional approach has a trade-off: JAX requires your functions to be pure — no side effects, no hidden state. We'll cover this requirement deeply in Track 3, but for now, understand that this purity constraint is what enables the composability. If a function modifies global state or mutates arrays, JAX can't safely reason about how to transform it.