Practical Example: Training Linear Regression from Scratch

Let's put everything together and train a linear regression model using pure JAX — no libraries, just grad, jit, and arrays. This demonstrates the complete JAX training pattern.

import jax
import jax.numpy as jnp

# ============================================
# 1. Generate synthetic data
# ============================================
key = jax.random.PRNGKey(0)
k1, k2, k3 = jax.random.split(key, 3)

n_samples, n_features = 500, 5
true_weights = jax.random.normal(k1, (n_features,))
true_bias = jnp.array(0.7)

X = jax.random.normal(k2, (n_samples, n_features))
noise = 0.1 * jax.random.normal(k3, (n_samples,))
y = X @ true_weights + true_bias + noise

# ============================================
# 2. Define model and loss (pure functions)
# ============================================
def predict(params, x):
    w, b = params
    return x @ w + b

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

# ============================================
# 3. Define training step (compiled)
# ============================================
@jax.jit
def train_step(params, x, y, lr):
    loss_val, grads = jax.value_and_grad(mse_loss)(params, x, y)
    w, b = params
    gw, gb = grads
    new_params = (w - lr * gw, b - lr * gb)
    return new_params, loss_val

# ============================================
# 4. Initialize and train
# ============================================
params = (jnp.zeros(n_features), jnp.array(0.0))
learning_rate = 0.1

for epoch in range(200):
    params, loss = train_step(params, X, y, learning_rate)
    if epoch % 40 == 0:
        print(f"Epoch {epoch:3d} | Loss: {loss:.6f}")

# ============================================
# 5. Compare results
# ============================================
learned_w, learned_b = params
print(f"\\nTrue weights:    {true_weights}")
print(f"Learned weights: {learned_w}")
print(f"True bias:       {true_bias}")
print(f"Learned bias:    {learned_b:.4f}")
print(f"Weight error:    {jnp.linalg.norm(true_weights - learned_w):.6f}")

Key observations about this training loop:

• Parameters are tuples: (w, b) — plain Python data structures, no special Parameter class.
• Gradients match parameter structure: jax.grad returns a (gw, gb) tuple matching the params tuple.
• Updates are functional: We create new_params instead of modifying params in place.
• Everything is compiled: The @jax.jit decorator compiles the entire train step including loss, gradients, and updates.

In production JAX code, you'd use Optax for the optimizer:

import optax

optimizer = optax.adam(learning_rate=1e-3)
opt_state = optimizer.init(params)

@jax.jit
def train_step_optax(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(mse_loss)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss