Intermediate

Neural Networks in JAX

JAX itself does not include a neural network library. Instead, you choose from ecosystem libraries like Flax (Google) or Haiku (DeepMind), and use Optax for optimization.

Flax: Google's Neural Network Library

Flax is the most popular neural network library for JAX. It uses an nn.Module system similar to PyTorch but with explicit parameter management:

Python

import jax
import jax.numpy as jnp
from flax import linen as nn

class MLP(nn.Module):
    hidden_dim: int = 128
    num_classes: int = 10

    @nn.compact
    def __call__(self, x, training: bool = False):
        x = nn.Dense(self.hidden_dim)(x)
        x = nn.relu(x)
        x = nn.Dropout(rate=0.3, deterministic=not training)(x)
        x = nn.Dense(self.hidden_dim)(x)
        x = nn.relu(x)
        x = nn.Dense(self.num_classes)(x)
        return x

# Initialize model and parameters
model = MLP()
key = jax.random.PRNGKey(0)
dummy_input = jnp.ones((1, 784))
params = model.init(key, dummy_input)

# Forward pass
logits = model.apply(params, x_batch)

Haiku: DeepMind's Alternative

Python

import haiku as hk

def mlp_fn(x):
    x = hk.Linear(128)(x)
    x = jax.nn.relu(x)
    x = hk.Linear(10)(x)
    return x

# Transform into pure functions
model = hk.transform(mlp_fn)

# Initialize parameters
key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, 784)))

# Forward pass
logits = model.apply(params, key, x_batch)

Training with Optax

Optax is JAX's gradient processing library. It provides optimizers, learning rate schedules, and gradient transformations:

Python

import optax

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

# Loss function
def loss_fn(params, x, y):
    logits = model.apply(params, x)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()

# Training step (JIT-compiled for speed)
@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

# Training loop
for epoch in range(20):
    for x_batch, y_batch in train_loader:
        params, opt_state, loss = train_step(params, opt_state, x_batch, y_batch)
    print(f"Epoch {epoch+1}, Loss: {loss:.4f}")

Flax vs Haiku

Feature	Flax (linen)	Haiku
Maintainer	Google	DeepMind
Style	Class-based (like PyTorch)	Function-based
Popularity	More popular, growing	Widely used at DeepMind
Training state	Built-in TrainState	Manual management
Recommendation	Best for most users	Good if you prefer functional style

Complete MNIST Example with Flax

Python

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.training import train_state
import optax

class CNN(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))
        x = nn.Dense(features=256)(x)
        x = nn.relu(x)
        x = nn.Dense(features=10)(x)
        return x

# Initialize
model = CNN()
key = jax.random.PRNGKey(0)
params = model.init(key, jnp.ones((1, 28, 28, 1)))

# Create TrainState (bundles params + optimizer)
state = train_state.TrainState.create(
    apply_fn=model.apply,
    params=params['params'],
    tx=optax.adam(1e-3)
)

Next Up: Advanced Topics

Learn multi-GPU parallelism with pmap, model sharding, custom gradients, and Hugging Face integration.

Next: Advanced Topics →

← Core Concepts Advanced Topics →

Ready to Go Deeper?

Live instructor-led courses from our partners. Affiliate disclosure.

🧠

AI & ML Courses - 30% Off

Live instructor-led AI, machine learning, data science, and cloud courses for working professionals. Use code Limited30 at checkout.

Edureka

📊

DataCamp - AI & Data Science

Hands-on Python, machine learning, and AI courses with interactive exercises and real projects.

DataCamp

🎓

edX - Top AI Courses

University-level AI courses from MIT, Harvard, Stanford. Earn certificates that employers recognize.

edX