Your First JAX Program: Putting It All Together

Let's write a complete JAX program that uses all four pillars — jit, grad, vmap, and basic array operations — to build a tiny linear regression from scratch. This will give you a feel for the JAX workflow before we dive deep into each component.

import jax
import jax.numpy as jnp

# 1. Generate synthetic data
key = jax.random.PRNGKey(42)
key, subkey = jax.random.split(key)
X = jax.random.normal(subkey, shape=(100, 3))  # 100 samples, 3 features

true_w = jnp.array([2.0, -1.0, 0.5])
y = X @ true_w + 0.1 * jax.random.normal(key, shape=(100,))

# 2. Define the model and loss as pure functions
def predict(params, x):
    return jnp.dot(x, params)

def loss_fn(params, X, y):
    preds = predict(params, X)
    return jnp.mean((preds - y) ** 2)

# 3. Compile the gradient computation
@jax.jit
def update(params, X, y, lr=0.1):
    grads = jax.grad(loss_fn)(params, X, y)
    return params - lr * grads  # Simple gradient descent

# 4. Train
params = jnp.zeros(3)  # Start from zeros
for step in range(100):
    params = update(params, X, y)
    if step % 20 == 0:
        current_loss = loss_fn(params, X, y)
        print(f"Step {step}, Loss: {current_loss:.4f}")

print(f"Learned params: {params}")
print(f"True params:    {true_w}")

Let's break down what makes this "JAX-style":

• No classes: The model is just a function (predict). Parameters are a plain array passed as an argument.
• Explicit state: Parameters aren't hidden inside an object. They're passed into and returned from functions.
• Functional updates: update doesn't modify params in place — it returns a new array.
• Composable transforms: jax.grad wraps loss_fn to get gradients, and jax.jit wraps the whole update for speed.
• Explicit randomness: Random numbers use explicit PRNG keys, not global state. We'll explore this in Track 3.

Now let's see vmap in action — computing per-example gradients, which is something unique to JAX's design:

# Per-example gradients: gradient of loss for EACH data point
# In PyTorch, this requires special tricks. In JAX, it's one line.

def single_loss(params, x, y):
    """Loss for a single example."""
    pred = jnp.dot(x, params)
    return (pred - y) ** 2

# vmap over the data dimensions (axis 0 of x and y), not params
per_example_grad_fn = jax.vmap(jax.grad(single_loss), in_axes=(None, 0, 0))

per_example_grads = per_example_grad_fn(params, X, y)
print(per_example_grads.shape)  # (100, 3) — one gradient per example

With just vmap(grad(f)), we computed 100 individual gradients without writing any loop or batch dimension logic. This is the power of composable transformations. In the following tracks, we'll go deep on each of these concepts.