Skip to content

JAX & Functional ML

Overview

JAX is a functional framework for ML: - Pure functional: All functions are pure (no side effects) - Automatic differentiation: jax.grad() for derivatives - JIT compilation: jax.jit() for speed - Vectorization: jax.vmap() for parallelism - Composable transformations: Chain transformations - Array operations: NumPy-like API with GPU support

JAX demonstrates how functional programming powers modern ML.

JAX Fundamentals

Installing JAX

# CPU version
pip install jax jaxlib

# GPU version (CUDA 11)
pip install --upgrade jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Apple Silicon
pip install --upgrade jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html

Pure Functions Required

import jax
import jax.numpy as jnp

# Pure function (works with JAX)
def pure_loss(params, x, y):
    """Pure loss function - no side effects."""
    predictions = jnp.dot(x, params)
    return jnp.mean((predictions - y) ** 2)

# Automatic differentiation
grad_fn = jax.grad(pure_loss)

# Compute gradients
params = jnp.array([1.0, 2.0, 3.0])
x = jnp.array([1.0, 2.0, 3.0])
y = jnp.array(14.0)

grads = grad_fn(params, x, y)
print(grads)

# IMPURE function (doesn't work with JAX)
loss_history = []

def impure_loss(params, x, y):
    """Impure - has side effects."""
    predictions = jnp.dot(x, params)
    loss = jnp.mean((predictions - y) ** 2)
    loss_history.append(float(loss))  # Side effect!
    return loss

# This doesn't work well with jax.grad()
# JAX expects pure functions!

jax.grad: Automatic Differentiation

Computing Derivatives

import jax
import jax.numpy as jnp

# Define function
def f(x):
    return x ** 3 + 2 * x ** 2 - 5 * x + 1

# Get derivative function
df = jax.grad(f)

# Compute derivative at x=2
derivative_at_2 = df(2.0)
print(f"f'(2) = {derivative_at_2}")

# Compute second derivative
d2f = jax.grad(df)
second_derivative = d2f(2.0)
print(f"f''(2) = {second_derivative}")

# Numerical verification
eps = 1e-5
numerical_grad = (f(2.0 + eps) - f(2.0 - eps)) / (2 * eps)
print(f"Numerical gradient: {numerical_grad}")

Multi-argument Differentiation

import jax
import jax.numpy as jnp

def loss(params, x, y):
    """Loss function with multiple arguments."""
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** 2)

# Differentiate with respect to first argument (params)
grad_params = jax.grad(loss)

# Differentiate with respect to specific arguments
grad_params_and_x = jax.grad(loss, argnums=(0, 1))

# Use it
params = jnp.array([1.0, 2.0])
x = jnp.array([1.0, 2.0])
y = 5.0

# Gradient w.r.t. params
grads = grad_params(params, x, y)
print(f"Gradient w.r.t. params: {grads}")

# Gradient w.r.t. both params and x
grads_both = grad_params_and_x(params, x, y)
print(f"Gradient w.r.t. params and x: {grads_both}")

jax.jit: Just-In-Time Compilation

Compile for Speed

import jax
import jax.numpy as jnp
import time

# Define function
def slow_function(x):
    for i in range(1000):
        x = x ** 2 + x
    return x

# Compile with jit
fast_function = jax.jit(slow_function)

# First call (includes compilation)
x = jnp.array(0.1)
start = time.time()
result1 = fast_function(x)
first_time = time.time() - start

# Subsequent calls (use compiled code)
start = time.time()
for _ in range(100):
    result = fast_function(x)
subsequent_time = time.time() - start

print(f"First call: {first_time*1000:.1f}ms (includes compilation)")
print(f"100 calls: {subsequent_time*1000:.1f}ms ({subsequent_time/100*1000:.3f}ms each)")

# Or use as decorator
@jax.jit
def compiled_function(x):
    return x ** 2 + x + 1

result = compiled_function(3.0)
print(result)

jax.vmap: Vectorization

Automatically Parallelize

import jax
import jax.numpy as jnp

def scalar_function(x):
    """Function that works on scalars."""
    return x ** 2 + jnp.sin(x)

# Vectorize to work on arrays
vectorized = jax.vmap(scalar_function)

# Works on entire array
x_array = jnp.array([1.0, 2.0, 3.0, 4.0])
results = vectorized(x_array)
print(results)  # [f(1), f(2), f(3), f(4)]

# Vectorize along specific axes
def scalar_matmul(row, matrix):
    """Multiply row by matrix."""
    return jnp.dot(row, matrix)

# Vectorize over first argument
vectorized_matmul = jax.vmap(scalar_matmul, in_axes=(0, None))

rows = jnp.array([[1.0, 2.0], [3.0, 4.0]])
matrix = jnp.array([[5.0, 6.0], [7.0, 8.0]])

results = vectorized_matmul(rows, matrix)
print(results)

Composing Transformations

Combining grad, jit, vmap

import jax
import jax.numpy as jnp

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

# Vectorize to compute loss for batch
def loss_batch(params, batch_x, batch_y):
    """Loss for batch of examples."""
    losses = jax.vmap(loss_single, in_axes=(None, 0, 0))(params, batch_x, batch_y)
    return jnp.mean(losses)

# Get gradient
grad_loss = jax.grad(loss_batch)

# JIT compile for speed
compiled_grad = jax.jit(grad_loss)

# Use in training
params = jnp.array([1.0, 2.0, 3.0])
batch_x = jnp.array([[1.0, 0.0, 0.0],
                      [0.0, 1.0, 0.0],
                      [0.0, 0.0, 1.0]])
batch_y = jnp.array([2.0, 4.0, 6.0])

grads = compiled_grad(params, batch_x, batch_y)
print(f"Gradients: {grads}")

# Training step
learning_rate = 0.01
new_params = params - learning_rate * grads
print(f"Updated params: {new_params}")

Functional Neural Networks

Stateless Models with JAX

import jax
import jax.numpy as jnp

# Functional neural network (no state)
def forward(params, x):
    """Forward pass with explicit parameters."""
    W1, b1, W2, b2 = params

    # Hidden layer
    hidden = jnp.dot(x, W1) + b1
    hidden = jnp.maximum(hidden, 0)  # ReLU

    # Output layer
    output = jnp.dot(hidden, W2) + b2
    return output

# Define loss
def loss_fn(params, batch):
    x, y = batch

    # Vectorize forward pass
    forward_batch = jax.vmap(lambda x: forward(params, x))
    predictions = forward_batch(x)

    return jnp.mean((predictions - y) ** 2)

# Initialize parameters
key = jax.random.PRNGKey(0)
params = [
    jax.random.normal(key, (10, 20)),  # W1
    jnp.zeros(20),                      # b1
    jax.random.normal(key, (20, 1)),   # W2
    jnp.zeros(1)                        # b2
]

# Training
grad_fn = jax.grad(loss_fn)
learning_rate = 0.01

for step in range(10):
    # Get gradients
    grads = grad_fn(params, (jnp.zeros((5, 10)), jnp.ones((5, 1))))

    # Update parameters
    params = [p - learning_rate * g for p, g in zip(params, grads)]

    print(f"Step {step}")

Common JAX Patterns

Pure Training Loop

import jax
import jax.numpy as jnp
from functools import partial

# Define model and loss separately
def model(params, x):
    """Pure model function."""
    return jnp.dot(x, params)

def loss(params, x, y):
    """Pure loss function."""
    predictions = model(params, x)
    return jnp.mean((predictions - y) ** 2)

# Get update function
@jax.jit
def update_step(params, x, y, learning_rate):
    """Single training step."""
    grads = jax.grad(loss)(params, x, y)
    return params - learning_rate * grads

# Training loop (purely functional)
def train(params, x, y, learning_rate, epochs):
    """Train model."""
    for epoch in range(epochs):
        params = update_step(params, x, y, learning_rate)
    return params

# Use it
params = jnp.array([1.0, 2.0, 3.0])
x = jnp.array([1.0, 0.0, 0.0](/1.0,-0.0,-0.0/))
y = jnp.array([2.0])

params = train(params, x, y, learning_rate=0.01, epochs=10)
print(f"Final params: {params}")

Advantages of Functional ML

Composability

# Chain transformations easily
@jax.jit
@partial(jax.vmap, in_axes=(None, 0, 0))
def batched_loss(params, x_batch, y_batch):
    return jax.grad(lambda p: loss(p, x_batch, y_batch))(params)

# Works seamlessly with jit and vmap together

Automatic Differentiation

# Compute any derivative you want
loss_val, loss_grad = jax.value_and_grad(loss)(params, x, y)

# Second derivatives
hessian = jax.hessian(loss)

# Directional derivatives
jvp = jax.jvp(loss, (params,), (direction,))

Hardware Agnostic

# Same code runs on CPU, GPU, TPU
import os
os.environ['JAX_PLATFORM_NAME'] = 'cpu'  # Use CPU
# or 'gpu' or 'tpu'

# Code unchanged, just runs on different hardware!

Summary: JAX vs PyTorch

Feature JAX PyTorch
Paradigm Functional Imperative
Transformations @jax.grad, @jax.jit, @jax.vmap Autograd, tracing
State Explicit (parameters passed) Implicit (object state)
Debugging Harder (functional) Easier (imperative)
Performance Excellent (JIT) Good (imperative)
Composability Excellent Good
Adoption Growing (research) Mature (production)