JAX Architecture: Design Patterns and Project Structure

JAX Architecture Overview

JAX is a high-performance numerical computing library that combines autograd capabilities with XLA (Accelerated Linear Algebra) compilation. Its architecture revolves around composable function transformations, pure functional programming patterns, and just-in-time compilation that targets CPUs, GPUs, and TPUs seamlessly. Understanding JAX's architectural philosophy is essential for building maintainable, performant, and scalable machine learning projects.

At its core, JAX provides four fundamental transformations that can be arbitrarily composed:

• grad – automatic differentiation for computing gradients
• jit – just-in-time compilation via XLA for hardware acceleration
• vmap – automatic vectorization for batching computations
• pmap – parallel mapping across multiple devices

These transformations operate on pure functions, meaning functions whose outputs depend solely on their inputs, with no side effects or mutable state. This constraint is not merely a stylistic preference — it is fundamental to how JAX guarantees correctness under arbitrary transformation composition.

Why Design Patterns Matter in JAX

Unlike frameworks such as PyTorch or TensorFlow that encourage object-oriented model definitions with mutable parameter storage, JAX demands a functional approach. This shift requires developers to adopt specific design patterns that:

• Maintain functional purity across the entire computation graph
• Enable seamless composition of transformations (e.g., jit(grad(fn)))
• Keep parameter management explicit and traceable
• Support reproducible research through stateless computation
• Allow efficient scaling from single-device prototyping to multi-host TPU pods

Without these patterns, JAX code quickly becomes tangled with hard-to-debug transformation conflicts, unexpected performance cliffs, and brittle parameter passing conventions.

Core Design Patterns

1. Functional Parameter Passing

The most fundamental pattern in JAX is passing all model parameters explicitly through function arguments rather than storing them in object attributes. This ensures functions remain pure and composable with transformations.

# ❌ Anti-pattern: mutable state in classes
class BadModel:
    def __init__(self):
        self.weights = jnp.array([1.0, 2.0])
    
    def predict(self, x):
        # Self is mutable state, breaks JAX transformations
        return x @ self.weights

# ✅ Correct pattern: explicit parameter passing
def predict(params, x):
    """Pure function: all inputs are explicit arguments."""
    w, b = params
    return x @ w + b

def loss_fn(params, x, y):
    """Composable loss function."""
    pred = predict(params, x)
    return jnp.mean((pred - y) ** 2)

# Now we can compose transformations freely
grad_loss = jax.grad(loss_fn)
value_and_grad = jax.value_and_grad(loss_fn)
compiled_fn = jax.jit(value_and_grad)

2. The Parameter-Update Loop Pattern

Training loops in JAX follow a functional update pattern where parameters are explicitly returned from update functions and carried forward to the next iteration. This is commonly implemented using jax.lax.scan for efficient compilation of entire training loops.

def train_step(state, batch):
    """Single training step: takes state, returns updated state."""
    params, opt_state, rng = state
    x, y = batch
    
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    
    # Update parameters functionally
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = jax.tree_util.tree_map(
        lambda p, u: p - 0.01 * u, params, updates
    )
    
    return (new_params, new_opt_state, rng), loss

# Compile the entire training loop
def train_epoch(params, opt_state, rng, data):
    """Process all batches using scan for maximum efficiency."""
    carry = (params, opt_state, rng)
    
    def step_fn(carry, batch):
        new_carry, loss = train_step(carry, batch)
        return new_carry, loss
    
    final_carry, losses = jax.lax.scan(step_fn, carry, data)
    return final_carry, losses

compiled_epoch = jax.jit(train_epoch)

3. The pytree Container Pattern

JAX uses "pytrees" — nested structures of arrays, dicts, lists, and tuples — as the universal container for parameters, optimizer states, and any other structured data. The jax.tree_util module provides utilities to manipulate these structures without breaking functional purity.

# Define complex parameter structures naturally
def init_params(key):
    k1, k2, k3 = jax.random.split(key, 3)
    return {
        'encoder': {
            'w1': jax.random.normal(k1, (784, 256)),
            'b1': jnp.zeros(256),
            'w2': jax.random.normal(k2, (256, 128)),
            'b2': jnp.zeros(128),
        },
        'decoder': {
            'w1': jax.random.normal(k3, (128, 784)),
            'b1': jnp.zeros(784),
        }
    }

# Use pytree operations for consistent updates
def add_l2_penalty(grads, params, lambda_=1e-4):
    """Add L2 penalty to gradients across entire pytree."""
    return jax.tree_util.tree_map(
        lambda g, p: g + lambda_ * p,
        grads,
        params
    )

# Flatten for serialization or analysis
def flatten_params(params):
    leaves, tree_def = jax.tree_util.tree_flatten(params)
    # leaves is a flat list of all arrays
    # tree_def stores the structure for reconstruction
    return leaves, tree_def

4. The Stochastic PRNGKey Pattern

JAX uses explicit random key management rather than global random state. Keys are threaded through functions and split deterministically, enabling full reproducibility even under parallel execution and JIT compilation.

def stochastic_layer(params, x, key):
    """Layer with dropout — key is explicit."""
    w, b = params
    keep_prob = 0.8
    # Split key: one for this op, one for downstream use
    dropout_key, next_key = jax.random.split(key)
    mask = jax.random.bernoulli(dropout_key, keep_prob, shape=x.shape)
    x = x * mask / keep_prob  # Scale to maintain expected value
    return x @ w + b, next_key  # Return new key for caller

# Chain multiple stochastic operations
def forward(params, x, key):
    k1, k2, k3 = jax.random.split(key, 3)
    h1, _ = stochastic_layer(params['layer1'], x, k1)
    h2, _ = stochastic_layer(params['layer2'], h1, k2)
    out, _ = stochastic_layer(params['layer3'], h2, k3)
    return out

# Reproducible initialization
init_key = jax.random.PRNGKey(42)
params = init_params(init_key)

5. The Transformation Composition Pattern

JAX's power comes from composing transformations. The order matters and follows specific rules. Understanding the transformation stack is crucial for correct and performant code.

# Building a full training pipeline through composition
def create_training_fn(model_fn, optimizer):
    """Compose transformations for a complete training function."""
    
    # Base loss function
    def loss_fn(params, batch, key):
        x, y = batch
        pred = model_fn(params, x, key)
        return jnp.mean(jnp.sum(-y * jax.nn.log_softmax(pred), axis=-1))
    
    # 1. Take gradient w.r.t. first argument (params)
    grad_fn = jax.grad(loss_fn)
    
    # 2. Vectorize over batch dimension if needed
    # grad_fn = jax.vmap(grad_fn, in_axes=(None, 0, None))
    
    # 3. JIT compile the gradient computation
    compiled_grad_fn = jax.jit(grad_fn)
    
    def update_fn(params, opt_state, batch, key):
        grads = compiled_grad_fn(params, batch, key)
        updates, new_opt_state = optimizer.update(grads, opt_state)
        new_params = jax.tree_util.tree_map(
            lambda p, u: p + u, params, updates
        )
        return new_params, new_opt_state
    
    return jax.jit(update_fn)

# The composition: jit(grad(loss_fn))
# This is the canonical JAX pattern

6. The Device-Agnostic Computation Pattern

JAX treats device placement as a transformation concern rather than a manual assignment problem. By default, computations run on the default device, but jax.device_put and jax.pmap handle distribution transparently.

# Single-device code that scales to multiple devices
def create_distributed_train_step(loss_fn, optimizer_update):
    """Same function works on 1 GPU or 8 GPUs."""
    
    def single_device_step(params, opt_state, batch, key):
        loss, grads = jax.value_and_grad(loss_fn)(params, batch, key)
        updates, new_opt_state = optimizer_update(grads, opt_state)
        new_params = jax.tree_util.tree_map(
            lambda p, u: p - 0.001 * u, params, updates
        )
        return (new_params, new_opt_state, key), loss
    
    # pmap transforms single-device fn to multi-device fn
    return jax.pmap(
        single_device_step,
        in_axes=(0, 0, 0, 0),  # Map over leading axis for each arg
        axis_name='devices'
    )

# Use the same pattern regardless of device count
step_fn = create_distributed_train_step(loss_fn, optimizer.update)
# step_fn now works with replicated inputs across all available devices

Project Structure Patterns

A well-structured JAX project separates concerns along functional boundaries rather than class hierarchies. Here is a battle-tested directory layout that scales from research experiments to production systems.

project/
├── src/
│   ├── __init__.py
│   ├── core/
│   │   ├── __init__.py
│   │   ├── transformations.py   # Custom JAX transformations
│   │   ├── pytree_utils.py      # pytree manipulation utilities
│   │   └── key_management.py    # PRNGKey helpers
│   ├── modeling/
│   │   ├── __init__.py
│   │   ├── modules.py           # Pure function model components
│   │   ├── attention.py         # Specific architectural patterns
│   │   └── embeddings.py
│   ├── training/
│   │   ├── __init__.py
│   │   ├── loops.py             # scan-based training loops
│   │   ├── losses.py            # Loss function library
│   │   └── metrics.py
│   ├── data/
│   │   ├── __init__.py
│   │   ├── pipeline.py          # tf.data or grain pipelines
│   │   └── augmentation.py
│   └── utils/
│       ├── __init__.py
│       ├── logging.py
│       ├── checkpointing.py     # orbax-based checkpointing
│       └── profiling.py
├── configs/
│   ├── default.py
│   └── experiment.py
├── experiments/
│   ├── train.py                 # Entry point scripts
│   └── evaluate.py
├── tests/
│   └── ...
├── setup.py
└── README.md

Module Design Principles

Every module in a JAX project exports pure functions, not classes with internal state. Here's how to structure a modeling module correctly:

# src/modeling/modules.py — correct functional module design

def linear_layer(params, x):
    """Pure linear transformation."""
    w, b = params
    return x @ w + b

def mlp(params, x):
    """Multi-layer perceptron — params is a nested pytree."""
    h = linear_layer(params['layer1'], x)
    h = jax.nn.relu(h)
    h = linear_layer(params['layer2'], h)
    h = jax.nn.relu(h)
    return linear_layer(params['output'], h)

def init_mlp(key, input_dim, hidden_dims, output_dim):
    """Initialize parameters for an MLP."""
    keys = jax.random.split(key, len(hidden_dims) + 2)
    dims = [input_dim] + hidden_dims + [output_dim]
    
    params = {}
    for i, (k, d_in, d_out) in enumerate(
        zip(keys, dims[:-1], dims[1:])
    ):
        w_key, b_key = jax.random.split(k)
        w = jax.random.normal(w_key, (d_in, d_out)) * jnp.sqrt(2.0 / d_in)
        b = jnp.zeros(d_out)
        params[f'layer{i}'] = (w, b)
    
    return params

# Export: functions only, no global state
__all__ = ['linear_layer', 'mlp', 'init_mlp']

Configuration Management Pattern

JAX projects benefit from explicit configuration objects passed as immutable structures. This maintains functional purity while providing flexibility.

# configs/default.py
import dataclasses
from typing import Tuple

@dataclasses.dataclass(frozen=True)
class ModelConfig:
    """Immutable configuration — safe to pass through transformations."""
    input_dim: int = 784
    hidden_dims: Tuple[int, ...] = (256, 128)
    output_dim: int = 10
    dropout_rate: float = 0.1
    
    def to_dict(self):
        return dataclasses.asdict(self)

@dataclasses.dataclass(frozen=True)
class TrainingConfig:
    learning_rate: float = 1e-3
    batch_size: int = 256
    num_epochs: int = 100
    seed: int = 42

# Usage in training script
def create_training_state(config):
    """Create initial training state from config."""
    key = jax.random.PRNGKey(config.seed)
    params = init_mlp(key, config.input_dim, 
                      list(config.hidden_dims), config.output_dim)
    opt_state = optimizer.init(params)
    return params, opt_state, key

Checkpointing with Orbax

Production JAX projects use Orbax for checkpointing, which integrates naturally with the pytree pattern.

# src/utils/checkpointing.py
import orbax.checkpoint as ocp

def save_checkpoint(path, step, params, opt_state, metadata=None):
    """Save training state as a checkpoint."""
    checkpoint_manager = ocp.CheckpointManager(
        path,
        ocp.PyTreeCheckpointHandler()
    )
    
    # Combine state into a single pytree
    state = {
        'params': params,
        'opt_state': opt_state,
        'step': step,
        'metadata': metadata or {}
    }
    
    checkpoint_manager.save(step, state)
    return state

def restore_checkpoint(path, step=None):
    """Restore training state from checkpoint."""
    checkpoint_manager = ocp.CheckpointManager(
        path,
        ocp.PyTreeCheckpointHandler()
    )
    
    if step is None:
        step = checkpoint_manager.latest_step()
    
    restored = checkpoint_manager.restore(step)
    return restored['params'], restored['opt_state'], step

Advanced Architectural Patterns

1. Custom Transformation Rules

When you need custom differentiation rules (e.g., for numerical stability), JAX allows defining custom gradient rules using jax.custom_jvp or jax.custom_vjp.

# Custom gradient for numerically stable log-sum-exp
@jax.custom_jvp
def stable_logsumexp(x, axis=-1):
    """Numerically stable log-sum-exp with custom gradient."""
    max_x = jnp.max(x, axis=axis, keepdims=True)
    shifted_x = x - max_x
    sum_exp = jnp.sum(jnp.exp(shifted_x), axis=axis)
    return max_x.squeeze(axis) + jnp.log(sum_exp)

@stable_logsumexp.defjvp
def stable_logsumexp_jvp(primals, tangents):
    """Custom Jacobian-vector product for efficiency."""
    x, axis = primals
    x_dot, _ = tangents
    
    max_x = jnp.max(x, axis=axis, keepdims=True)
    shifted_x = x - max_x
    exp_x = jnp.exp(shifted_x)
    sum_exp = jnp.sum(exp_x, axis=axis)
    
    # Stable softmax for gradient
    softmax = exp_x / sum_exp
    primal_out = stable_logsumexp(x, axis)
    tangent_out = jnp.sum(softmax * x_dot, axis=axis)
    
    return primal_out, tangent_out

2. Conditional Computation with JIT

JIT compilation requires static shapes. Dynamic control flow must use JAX's structured control flow primitives like jax.lax.cond, jax.lax.scan, and jax.lax.while_loop.

def dynamic_forward(params, x, training_mode):
    """Conditional computation that works under JIT."""
    
    def training_path(p, x):
        # Apply dropout, batch norm in training mode
        x = apply_dropout(p['dropout'], x, training=True)
        return apply_batch_norm(p['bn'], x, training=True)
    
    def inference_path(p, x):
        # Skip dropout, use running stats for batch norm
        return apply_batch_norm(p['bn'], x, training=False)
    
    # Use lax.cond for JIT-compatible branching
    # Both branches must return same-shaped outputs
    return jax.lax.cond(
        training_mode,
        lambda: training_path(params, x),
        lambda: inference_path(params, x)
    )

# Scan for loops with dynamic length
def process_variable_length_sequences(params, sequences, lengths):
    """Process sequences of different lengths under JIT."""
    
    def step_fn(carry, _):
        state, idx = carry
        # Pad to max length, mask based on valid lengths
        valid_mask = idx < lengths
        # Apply computation with masking
        new_state = masked_computation(params, state, valid_mask)
        return (new_state, idx + 1), None
    
    max_len = sequences.shape[1]
    init_carry = (initial_state, 0)
    final_carry, _ = jax.lax.scan(step_fn, init_carry, None, length=max_len)
    return final_carry[0]

3. Multi-Device Orchestration

For multi-host TPU or multi-GPU training, JAX uses pmap with explicit axis naming for collectives and sharding strategies.

def create_multi_device_train_step(loss_fn, optimizer):
    """Create a training step that works across multiple devices."""
    
    def step_fn(params, opt_state, batch, key):
        loss, grads = jax.value_and_grad(loss_fn)(params, batch, key)
        
        # Average gradients across devices using psum
        grads = jax.lax.pmean(
            grads, axis_name='devices'
        )
        
        updates, new_opt_state = optimizer.update(grads, opt_state)
        new_params = jax.tree_util.tree_map(
            lambda p, u: p - 1e-3 * u, params, updates
        )
        
        # Synchronize loss for logging
        loss = jax.lax.pmean(loss, axis_name='devices')
        
        return (new_params, new_opt_state, key), loss
    
    # pmap replicates computation across devices
    return jax.pmap(
        step_fn,
        in_axes=(0, 0, 0, 0),  # Replicate along device axis
        axis_name='devices'
    )

# Data sharding pattern
def shard_data(data, num_devices):
    """Split data across devices for pmap."""
    # Reshape to (num_devices, batch_per_device, ...)
    per_device_batch = data.shape[0] // num_devices
    return data.reshape(
        (num_devices, per_device_batch) + data.shape[1:]
    )

Testing Patterns for JAX Projects

Testing pure functions is straightforward, but JAX transformations introduce unique testing requirements.

def test_transformation_composition():
    """Verify that transformations compose correctly."""
    key = jax.random.PRNGKey(0)
    params = init_mlp(key, 4, [8, 8], 2)
    x = jnp.ones((16, 4))
    y = jnp.zeros((16, 2))
    
    def loss(p, x, y):
        return jnp.mean((mlp(p, x) - y) ** 2)
    
    # Test 1: grad works
    grads = jax.grad(loss)(params, x, y)
    assert all(g is not None for g in jax.tree_util.tree_leaves(grads))
    
    # Test 2: jit(grad) matches grad
    jit_grads = jax.jit(jax.grad(loss))(params, x, y)
    for g1, g2 in zip(
        jax.tree_util.tree_leaves(grads),
        jax.tree_util.tree_leaves(jit_grads)
    ):
        assert jnp.allclose(g1, g2, atol=1e-6)
    
    # Test 3: vmap(grad) works
    xs = jnp.ones((32, 16, 4))
    ys = jnp.zeros((32, 16, 2))
    vmap_grads = jax.vmap(
        jax.grad(loss), in_axes=(None, 0, 0)
    )(params, xs, ys)
    assert vmap_grads is not None
    
    # Test 4: Determinism under same key
    k1 = jax.random.PRNGKey(42)
    k2 = jax.random.PRNGKey(42)
    p1 = init_mlp(k1, 4, [8], 2)
    p2 = init_mlp(k2, 4, [8], 2)
    for l1, l2 in zip(
        jax.tree_util.tree_leaves(p1),
        jax.tree_util.tree_leaves(p2)
    ):
        assert jnp.allclose(l1, l2)

Performance Optimization Patterns

1. JIT Placement Strategy

Deciding where to place JIT boundaries significantly impacts both compilation time and runtime performance. The key insight is to JIT the largest possible computational unit without introducing recompilation.

# Good: JIT the entire training step
@jax.jit
def train_step(params, opt_state, batch, key):
    # All computations fused into one XLA graph
    loss, grads = jax.value_and_grad(loss_fn)(params, batch, key)
    updates, new_opt_state = optimizer.update(grads, opt_state)
    new_params = jax.tree_util.tree_map(
        lambda p, u: p - 1e-3 * u, params, updates
    )
    return (new_params, new_opt_state, key), loss

# Bad: JIT only small pieces — prevents fusion
def train_step_bad(params, opt_state, batch, key):
    loss, grads = jax.jit(jax.value_and_grad(loss_fn))(params, batch, key)
    # XLA cannot fuse across this boundary
    updates, new_opt_state = optimizer.update(grads, opt_state)
    return (new_params, new_opt_state, key), loss

2. Avoiding Recompilation

JIT recompilation happens when function signatures change. Keep static arguments (shapes, dtypes) consistent across calls.

# Use static_argnums for arguments that affect shapes
@functools.partial(jax.jit, static_argnums=(2,))
def compute_with_varying_shape(params, x, multiplier):
    """multiplier changes output shape — mark as static."""
    # multiplier determines output dimension
    output_dim = x.shape[-1] * multiplier
    w = params['projection']
    return x @ w[:, :output_dim]

# For truly dynamic shapes, pad to maximum and mask
@jax.jit
def compute_fixed_shape(params, x, valid_lengths):
    """Fixed max shape, use masking for variable lengths."""
    max_len = x.shape[1]
    mask = jnp.arange(max_len)[None, :] < valid_lengths[:, None]
    result = process_sequence(params, x)
    return result * mask[..., None]  # Zero out padding positions

3. Memory Management with Donated Buffers

For large models, use JAX's buffer donation to avoid unnecessary copies during parameter updates.

# Tell JAX that params can be reused in-place
@jax.jit
def update_params(params, grads, learning_rate):
    # Compute new params
    new_params = jax.tree_util.tree_map(
        lambda p, g: p - learning_rate * g,
        params,
        grads
    )
    return new_params

# Use with donated buffers
def train_loop(params, data):
    for batch in data:
        grads = jax.grad(loss_fn)(params, batch)
        # params buffer will be donated (reused) for output
        params = jax.jit(
            update_params,
            donate_argnums=(0,)  # params buffer can be reused
        )(params, grads, 0.01)
    return params

Common Pitfalls and Their Resolutions

Pitfall 1: Accidental Side Effects

# ❌ Dangerous: modifies array in-place under JIT
@jax.jit
def bad_update(params, grads):
    params[0][:] = params[0] - 0.01 * grads[0]  # Runtime error!
    return params

# ✅ Correct: always create new arrays
@jax.jit
def good_update(params, grads):
    return jax.tree_util.tree_map(
        lambda p, g: p - 0.01 * g,
        params,
        grads
    )

Pitfall 2: Python Control Flow Under JIT

# ❌ Python if/for don't work inside JIT with traced values
@jax.jit
def bad_loop(x):
    for i in range(x.shape[0]):  # x is traced, shape is dynamic
        x = x.at[i].set(x[i] * 2)
    return x

# ✅ Use lax.scan or vectorized operations
@jax.jit
def good_loop(x):
    return x * 2  # Vectorized — no Python loop needed

# Or for truly iterative algorithms:
@jax.jit
def good_scan(x):
    def body_fn(carry, _):
        return carry * 2, None
    final, _ = jax.lax.scan(body_fn, x, None, length=10)
    return final

Pitfall 3: Forgetting to Split Keys

# ❌ Reusing the same key — produces identical randomness
def bad_dropout(key, x):
    mask1 = jax.random.bernoulli(key, 0.5, x.shape)
    mask2 = jax.random.bernoulli(key, 0.5, x.shape)  # Same mask!
    return x * mask1 * mask2

# ✅ Always split keys before use
def good_dropout(key, x):
    k1, k2 = jax.random.split(key)
    mask1 = jax.random.bernoulli(k1, 0.5, x.shape)
    mask2 = jax.random.bernoulli(k2, 0.5, x.shape)
    return x * mask1 * mask2

Conclusion

JAX's architecture represents a fundamental departure from object-oriented deep learning frameworks. Its design patterns — functional parameter passing, explicit PRNGKey management, pytree containers, transformation composition, and structured control flow — form a coherent system that unlocks powerful capabilities. When structuring a JAX project, the guiding principle is to build libraries of pure functions organized by domain concern rather than by class hierarchy, keeping all state explicit and all side effects eliminated. This discipline pays dividends in performance through XLA compilation, in correctness through transformation composability, and in scalability through seamless device parallelism. By internalizing these patterns, developers can write JAX code that is not only fast and correct but also elegantly maintainable, reproducible, and ready to scale from a single GPU to a full TPU pod without architectural changes.