Practical Example: Porting NumPy to JAX

Let's walk through porting a real numerical computation from NumPy to JAX — a softmax function with a small neural network forward pass. This exercise highlights the typical changes you'll need to make.

Original NumPy Code

import numpy as np

def softmax_numpy(x):
    """Numerically stable softmax."""
    x_max = np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(x - x_max)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

def forward_numpy(params, x):
    """Two-layer neural network."""
    w1, b1, w2, b2 = params
    # Layer 1
    h = np.dot(x, w1) + b1
    # ReLU — uses in-place clipping in idiomatic NumPy
    h = np.maximum(h, 0)
    # Layer 2
    logits = np.dot(h, w2) + b2
    return softmax_numpy(logits)

# Initialize
np.random.seed(42)
w1 = np.random.randn(4, 8) * 0.1
b1 = np.zeros(8)
w2 = np.random.randn(8, 3) * 0.1
b2 = np.zeros(3)
params = [w1, b1, w2, b2]

x = np.random.randn(16, 4)  # 16 samples, 4 features
probs = forward_numpy(params, x)
print(probs.shape)  # (16, 3)
print(probs.sum(axis=-1))  # All 1.0

JAX Port — Minimal Changes

import jax
import jax.numpy as jnp

def softmax_jax(x):
    """Numerically stable softmax — identical logic."""
    x_max = jnp.max(x, axis=-1, keepdims=True)
    exp_x = jnp.exp(x - x_max)
    return exp_x / jnp.sum(exp_x, axis=-1, keepdims=True)

def forward_jax(params, x):
    """Two-layer neural network — pure function."""
    w1, b1, w2, b2 = params
    h = jnp.dot(x, w1) + b1
    h = jnp.maximum(h, 0)
    logits = jnp.dot(h, w2) + b2
    return softmax_jax(logits)

# Initialize with JAX's explicit PRNG
key = jax.random.PRNGKey(42)
keys = jax.random.split(key, 4)
w1 = jax.random.normal(keys[0], (4, 8)) * 0.1
b1 = jnp.zeros(8)
w2 = jax.random.normal(keys[1], (8, 3)) * 0.1
b2 = jnp.zeros(3)
params = (w1, b1, w2, b2)  # Tuple, not list (for pytree compatibility)

x = jax.random.normal(keys[2], (16, 4))
probs = forward_jax(params, x)
print(probs.shape)  # (16, 3)

What Changed?

1. np → jnp: All NumPy calls replaced with JAX NumPy.
2. Random numbers: np.random.seed(42) replaced with explicit jax.random.PRNGKey(42) and jax.random.split.
3. List → Tuple: Changed params from a list to a tuple. JAX's pytree system handles tuples natively for differentiation.
4. No other changes needed! The core computation logic is identical.

Now Add JAX Superpowers

# JIT compile for speed
fast_forward = jax.jit(forward_jax)
probs = fast_forward(params, x)  # First call compiles; subsequent calls are fast

# Get gradients of a loss function
def loss_fn(params, x, targets):
    probs = forward_jax(params, x)
    # Cross-entropy loss
    return -jnp.mean(jnp.sum(targets * jnp.log(probs + 1e-8), axis=-1))

# One-hot targets
targets = jax.nn.one_hot(jnp.array([0, 1, 2, 0, 1, 2, 0, 1,
                                      2, 0, 1, 2, 0, 1, 2, 0]), 3)

# Gradient with respect to params
grads = jax.grad(loss_fn)(params, x, targets)
print(type(grads))  # tuple — same structure as params!
print(grads[0].shape)  # (4, 8) — gradient for w1