← Back to DevBytes
Developer Tutorial July 17, 2026 14 min read

Migrating from PyTorch to JAX: Step-by-Step Guide

What is JAX and Why Migrate?

JAX is a high-performance numerical computing library from Google Research that provides composable transformations of Python+NumPy programs: automatic differentiation (grad), JIT compilation (jit), automatic vectorization (vmap), and parallelization (pmap). Unlike PyTorch's eager-first imperative style, JAX is built on a functional programming paradigm where computations are pure functions over immutable arrays.

Migrating from PyTorch to JAX matters for several compelling reasons:

Core Mental Model Shift

🚀 Deploy your AI agent in 10 minutes

Managed Hermes hosting. Zero DevOps. 100M tokens/mo included.

Try it free →

Before diving into code, understand the fundamental paradigm differences:

PyTorch: Imperative, Stateful, Object-Oriented

# PyTorch - everything is an object with mutable state
model = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters())
loss_fn = nn.CrossEntropyLoss()

def train_step(x, y):
    optimizer.zero_grad()        # mutate optimizer state
    output = model(x)            # model holds parameters internally
    loss = loss_fn(output, y)
    loss.backward()              # mutate .grad attributes
    optimizer.step()             # mutate parameters in-place
    return loss.item()

JAX: Functional, Stateless, Explicit

# JAX - everything is a pure function of explicit arguments
import jax
import jax.numpy as jnp
import optax

def forward(params, x):
    w, b = params
    return jnp.dot(x, w) + b     # pure function, no internal state

def loss_fn(params, x, y):
    pred = forward(params, x)
    return -jnp.mean(jnp.log(pred) * y)  # returns value, no side effects

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

def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state, loss

The key insight: everything becomes explicit arguments and return values. There are no hidden states, no in-place mutations, no object hierarchies holding internal variables.

Step-by-Step Migration Guide

Step 1: Environment Setup

Install JAX with the appropriate accelerator support:

# CPU-only version
pip install jax jaxlib

# CUDA GPU version (CUDA 12)
pip install "jax[cuda12]"

# TPU version
pip install "jax[tpu]"

Also install the ecosystem libraries you'll need:

pip install optax flax orbax-checkpoint

Step 2: Tensors and Array Operations

Replace torch.* tensor operations with jax.numpy (which mirrors NumPy but operates on immutable JAX arrays):

# PyTorch
import torch
x = torch.randn(32, 10)
y = torch.nn.functional.relu(x)
z = torch.matmul(x, x.T)

# JAX equivalent
import jax.numpy as jnp
from jax import random

key = random.PRNGKey(42)
x = random.normal(key, (32, 10))
y = jax.nn.relu(x)           # or: jnp.maximum(x, 0)
z = jnp.matmul(x, x.T)       # or: x @ x.T

Critical difference: JAX arrays are immutable. You cannot do x[0, 0] = 5. Instead, use index update syntax:

# Instead of x[0, 0] = 5 (in-place mutation)
x = x.at[0, 0].set(5)         # returns a NEW array
x = x.at[1:3].add(1.0)        # adds 1.0 to slice, returns new array

Step 3: Random Number Generation

This is perhaps the most jarring change. PyTorch uses a global mutable RNG state; JAX requires explicit PRNG keys that you split and pass around functionally:

# PyTorch - implicit global RNG state
x = torch.randn(100, 10)      # global state mutated invisibly
y = torch.randn(100, 10)      # different numbers, same global mutation

# JAX - explicit PRNG key management
from jax import random

key = random.PRNGKey(42)      # initial seed
key1, key2 = random.split(key)  # split for independent streams

x = random.normal(key1, (100, 10))
y = random.normal(key2, (100, 10))

# Pattern for training loops: split as you go
key, subkey = random.split(key)
dropout_mask = random.bernoulli(subkey, 0.5, (64, 128))
key, subkey = random.split(key)   # split again for next use
# ... continue splitting through the loop

The rule: never reuse a key. Every stochastic operation consumes a key and you split to get fresh ones. This is tedious at first but gives you exact reproducibility and parallel-safe randomness.

Step 4: Model Definition — Replacing nn.Module

PyTorch models are classes holding state. JAX models are typically either pure functions or Flax modules. Here's a side-by-side comparison of a simple MLP:

# ===== PYTHORCH =====
class MLP(nn.Module):
    def __init__(self, in_dim=784, hidden=256, out_dim=10):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.fc3 = nn.Linear(hidden, out_dim)
        self.dropout = nn.Dropout(0.1)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = MLP()
# Parameters are implicitly owned by model.parameters()

# ===== JAX WITH FLAX =====
import flax.linen as nn

class MLP(nn.Module):
    hidden: int = 256
    out_dim: int = 10
    dropout_rate: float = 0.1
    
    @nn.compact
    def __call__(self, x, deterministic=False):
        x = nn.Dense(self.hidden)(x)
        x = nn.relu(x)
        x = nn.Dropout(self.dropout_rate)(
            x, deterministic=deterministic)
        x = nn.Dense(self.hidden)(x)
        x = nn.relu(x)
        x = nn.Dense(self.out_dim)(x)
        return x

model = MLP()
# Parameters are initialized separately and passed explicitly
key = random.PRNGKey(0)
dummy_input = jnp.ones((1, 784))
variables = model.init(key, dummy_input)
params = variables['params']  # pure pytree of arrays

Key differences in Flax:

Step 5: Loss Functions and Automatic Differentiation

Replace loss.backward() and torch.no_grad() with JAX's functional gradient transforms:

# PyTorch
def train_step(model, optimizer, x, y):
    optimizer.zero_grad()
    loss = F.cross_entropy(model(x), y)
    loss.backward()
    optimizer.step()
    return loss.item()

# JAX
def compute_loss(params, x, y):
    logits = model.apply({'params': params}, x)
    return optax.softmax_cross_entropy_with_integer_labels(
        logits, y).mean()

# jax.grad returns a function that computes gradients
grad_fn = jax.grad(compute_loss)
# jax.value_and_grad returns both value and gradients
value_and_grad_fn = jax.value_and_grad(compute_loss)

def train_step(params, opt_state, x, y):
    loss, grads = value_and_grad_fn(params, x, y)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state, loss

JAX provides jax.grad for scalar-output functions, jax.jacfwd/jax.jacrev for full Jacobians, and jax.vjp/jax.jvp for vector-Jacobian products. These compose with other transformations:

# Compute per-example gradients efficiently with vmap+grad
per_example_grads = jax.vmap(
    jax.grad(compute_loss, argnums=0), 
    in_axes=(None, 0, 0)
)(params, x_batch, y_batch)

Step 6: Optimizers with Optax

Optax replaces torch.optim. It treats optimizers as pure functions that transform gradients and maintain state explicitly:

# PyTorch
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# optimizer holds internal state, mutates parameters in-place

# JAX with Optax
import optax

optimizer = optax.adam(learning_rate=1e-3)
# Or chain transformations:
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=1e-3, weight_decay=1e-4),
)

# Initialize optimizer state
opt_state = optimizer.init(params)

# In training loop:
grads = jax.grad(loss_fn)(params, x, y)
updates, new_opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)
# params, opt_state are now updated for the next iteration

Optax's chaining API is powerful — you can stack gradient clipping, weight decay, learning rate schedules, and the core optimizer into a single pipeline.

Step 7: Training Loop — Putting It All Together

Here's a complete training loop that demonstrates the full JAX pattern:

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

# --- Model definition ---
class CNN(nn.Module):
    num_classes: int = 10
    
    @nn.compact
    def __call__(self, x, deterministic=False):
        x = nn.Conv(features=32, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(features=64, kernel_size=(3, 3))(x)
        x = nn.relu(x)
        x = x.reshape(x.shape[0], -1)  # flatten
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(self.num_classes)(x)
        return x

# --- Loss function ---
def loss_fn(params, batch, deterministic=False):
    images, labels = batch
    logits = model.apply({'params': params}, images, 
                         deterministic=deterministic)
    loss = optax.softmax_cross_entropy_with_integer_labels(
        logits, labels).mean()
    return loss

# --- Metrics ---
def compute_accuracy(params, batch):
    images, labels = batch
    logits = model.apply({'params': params}, images, deterministic=True)
    preds = jnp.argmax(logits, axis=-1)
    return jnp.mean(preds == labels)

# --- Training step (to be JIT-compiled) ---
@jax.jit
def train_step(params, opt_state, batch):
    loss, grads = jax.value_and_grad(loss_fn)(params, batch)
    updates, new_opt_state = optimizer.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    return new_params, new_opt_state, loss

# --- Initialization ---
key = random.PRNGKey(42)
model = CNN(num_classes=10)

# Initialize parameters
key, init_key = random.split(key)
dummy_input = jnp.ones((1, 28, 28, 1))  # MNIST-like
variables = model.init(init_key, dummy_input)
params = variables['params']

# Initialize optimizer
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=1e-3, weight_decay=1e-4)
)
opt_state = optimizer.init(params)

# --- Training loop ---
num_epochs = 10
for epoch in range(num_epochs):
    # Dummy data loader
    for batch_idx in range(len(train_images) // 32):
        batch_images = train_images[batch_idx*32:(batch_idx+1)*32]
        batch_labels = train_labels[batch_idx*32:(batch_idx+1)*32]
        batch = (batch_images, batch_labels)
        
        params, opt_state, loss = train_step(params, opt_state, batch)
        
    # Evaluation (no JIT needed if small, but you can JIT this too)
    eval_batch = (test_images[:1000], test_labels[:1000])
    acc = compute_accuracy(params, eval_batch)
    print(f"Epoch {epoch}: loss={loss:.4f}, accuracy={acc:.4f}")

Step 8: JIT Compilation — The Performance Magic

The @jax.jit decorator compiles your function to XLA-optimized kernels. This is where massive speedups come from. Key rules for JIT-compatible code:

# This BREAKS JIT — data-dependent Python if
@jax.jit
def broken_fn(x):
    if x[0] > 0:        # x is an array, truth value is ambiguous
        return x
    return -x

# This WORKS — use jax.lax.cond for control flow
@jax.jit
def fixed_fn(x):
    return jax.lax.cond(x[0] > 0, lambda: x, lambda: -x)

# This BREAKS JIT — dynamic shape
@jax.jit
def broken_dynamic(x):
    return x[:random.randint(0, 10)]  # JIT needs static bounds

# This WORKS — use static argnames or fixed slices
@jax.jit
def fixed_slice(x):
    return x[:10]  # static slice

You can debug JIT issues by calling the function once without @jit to see the Python-side logic, then adding JIT gradually.

Step 9: Vectorization with vmap

vmap automatically vectorizes a function over a batch dimension, replacing manual batching code:

# PyTorch — batch processing is implicit in nn.Module
output = model(batch_images)  # automatically handles batches

# JAX — explicit vmap for per-sample operations
# Say you have a function that processes ONE image
def process_single(params, image):
    return model.apply({'params': params}, image[None, ...])[0]

# vmap creates a batched version automatically
process_batch = jax.vmap(process_single, in_axes=(None, 0))
batch_output = process_batch(params, batch_images)

# Even more powerful: vmap over multiple axes
# Compute per-example gradients for an entire batch
per_example_grad_fn = jax.vmap(
    jax.value_and_grad(loss_fn, argnums=0),
    in_axes=(None, (0, 0))  # params fixed, batch split
)
losses, per_example_grads = per_example_grad_fn(params, batch)

Step 10: Handling State — BatchNorm, Dropout, Running Stats

PyTorch uses model.train() / model.eval() and internal buffers. JAX requires explicit state passing:

# PyTorch — state is implicit
model.train()
out = model(x)         # updates running_mean internally
model.eval()
out = model(x)         # uses stored running_mean

# JAX with Flax — explicit state variables
class StatefulModule(nn.Module):
    @nn.compact
    def __call__(self, x, deterministic=False):
        x = nn.Dense(256)(x)
        x = nn.BatchNorm(use_running_average=not deterministic)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

# Initialize returns both params and batch_stats
variables = model.init(key, x)
params = variables['params']
batch_stats = variables['batch_stats']

# Training: pass mutable state, update it
def train_step(params, batch_stats, opt_state, batch):
    # mutable=['batch_stats'] tells Flax to update them
    (loss, logits), updated_vars = model.apply(
        {'params': params, 'batch_stats': batch_stats},
        batch['images'], deterministic=False,
        mutable=['batch_stats']
    )
    new_batch_stats = updated_vars['batch_stats']
    # Compute gradients on loss...
    return new_params, new_batch_stats, new_opt_state, loss

# Evaluation: use stored stats, don't update
def eval_step(params, batch_stats, batch):
    logits = model.apply(
        {'params': params, 'batch_stats': batch_stats},
        batch['images'], deterministic=True
    )
    return logits

This explicit state threading is more verbose but gives you total control — you can easily checkpoint, restore, or manipulate running statistics.

Step 11: Data Loading

JAX doesn't have a direct replacement for torch.utils.data.DataLoader. Common approaches:

# Option 1: PyTorch DataLoader bridge
from torch.utils.data import DataLoader
import numpy as np

def numpy_collate(batch):
    """Convert PyTorch tensors to numpy arrays"""
    return {k: v.numpy() if isinstance(v, torch.Tensor) else v 
            for k, v in batch.items()}

dataloader = DataLoader(dataset, batch_size=32, 
                        collate_fn=numpy_collate)

for batch in dataloader:
    images = jnp.asarray(batch['images'])  # zero-copy if possible
    labels = jnp.asarray(batch['labels'])
    # ... training step

# Option 2: Pure NumPy loader (simple datasets)
def numpy_data_stream(x_path, y_path, batch_size, key):
    x_data = np.load(x_path)
    y_data = np.load(y_path)
    num_samples = len(x_data)
    while True:
        key, subkey = random.split(key)
        indices = random.permutation(subkey, num_samples)
        for i in range(0, num_samples, batch_size):
            batch_idx = indices[i:i+batch_size]
            yield jnp.asarray(x_data[batch_idx]), jnp.asarray(y_data[batch_idx])

Step 12: Checkpointing and Model Persistence

Replace torch.save / torch.load with Orbax or simple pickling of pytrees:

# PyTorch
torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': epoch,
}, 'checkpoint.pt')
checkpoint = torch.load('checkpoint.pt')
model.load_state_dict(checkpoint['model_state_dict'])

# JAX with Orbax
import orbax.checkpoint as ocp

checkpointer = ocp.StandardCheckpointer()
checkpointer.save('checkpoint_dir/', 
    args=ocp.args.StandardSave(
        {'params': params, 'opt_state': opt_state, 'epoch': epoch}
    ))

# Restore
restored = checkpointer.restore('checkpoint_dir/',
    args=ocp.args.StandardRestore(
        {'params': params, 'opt_state': opt_state, 'epoch': 0}
    ))
params = restored['params']
opt_state = restored['opt_state']

# Simple alternative: pickle pytrees
import pickle
with open('checkpoint.pkl', 'wb') as f:
    pickle.dump({'params': params, 'opt_state': opt_state}, f)
with open('checkpoint.pkl', 'rb') as f:
    restored = pickle.load(f)

Common Migration Pitfalls and Best Practices

Pitfall 1: In-Place Mutation Habits

# WRONG: Trying to mutate JAX arrays
params['Dense_0']['kernel'][0, :] = 0.0  # ERROR!

# CORRECT: Use .at[] updates
params['Dense_0']['kernel'] = params['Dense_0']['kernel'].at[0, :].set(0.0)

Pitfall 2: Forgetting to Split PRNG Keys

# WRONG: Reusing the same key
key = random.PRNGKey(42)
x = random.normal(key, (100,))
y = random.normal(key, (100,))  # y will be IDENTICAL to x!

# CORRECT: Split before each use
key = random.PRNGKey(42)
key1, key2 = random.split(key)
x = random.normal(key1, (100,))
y = random.normal(key2, (100,))  # independent random values

Pitfall 3: Python Control Flow in JIT Functions

# WRONG: Data-dependent if/for inside @jit
@jax.jit
def bad_loop(x):
    for i in range(x.shape[0]):  # x.shape[0] is a traced value
        x = x.at[i].set(x[i] * 2)
    return x

# CORRECT: Use jax.lax.scan or static iteration
@jax.jit
def good_loop(x):
    return jax.lax.fori_loop(0, x.shape[0], 
        lambda i, arr: arr.at[i].set(arr[i] * 2), x)

Pitfall 4: Ignoring JIT Warm-Up

The first call to a @jax.jit function compiles it, which can take seconds. Always "warm up" JIT functions before benchmarking:

# Warm up JIT
_ = train_step(params, opt_state, dummy_batch)
# Now benchmark the real run
start = time.time()
for _ in range(100):
    params, opt_state, loss = train_step(params, opt_state, real_batch)
elapsed = time.time() - start

Pitfall 5: Expecting PyTorch-Style Debugging

JIT-compiled code doesn't support print() or pdb inside the function. Use:

Best Practice: Gradual JIT Adoption

Start with eager mode (no JIT) to verify correctness, then add JIT incrementally:

# Phase 1: Write and debug eagerly
def train_step(params, opt_state, batch):
    ...
# Phase 2: Add JIT once logic is correct
train_step = jax.jit(train_step, static_argnames=['deterministic'])

Best Practice: Use Pytrees Effectively

JAX's pytree abstraction (nested dicts/lists/tuples of arrays) is how you manage structured parameters. Libraries like Flax produce pytrees naturally:

# Flatten and unflatten for custom manipulation
from jax.tree_util import tree_map, tree_flatten

# Apply learning rate scaling to all parameters
def scale_lr(param):
    return param * 0.1
scaled_params = tree_map(scale_lr, params)

# Compute global norm across all parameters
leaves, _ = tree_flatten(params)
global_norm = jnp.sqrt(sum(jnp.sum(jnp.square(leaf)) for leaf in leaves))

Best Practice: Profile with jax.profiler

# Profile JIT-compiled code
with jax.profiler.trace('logs/'):
    for step in range(100):
        params, opt_state, loss = train_step(params, opt_state, batch)
# View with TensorBoard: tensorboard --logdir logs/

Full End-to-End Migration Example

Below is a complete, runnable training script that mirrors a typical PyTorch workflow — data loading, model definition, training loop, evaluation, and checkpointing — all in JAX:

import jax
import jax.numpy as jnp
import optax
import flax.linen as nn
from flax.training import train_state, checkpoints
from jax import random
import numpy as np

# ============================================================
# 1. Model Definition (replaces nn.Module)
# ============================================================
class ResBlock(nn.Module):
    features: int
    
    @nn.compact
    def __call__(self, x, deterministic=False):
        residual = x
        x = nn.Conv(self.features, kernel_size=(3, 3), padding='SAME')(x)
        x = nn.BatchNorm(use_running_average=not deterministic)(x)
        x = nn.relu(x)
        x = nn.Conv(self.features, kernel_size=(3, 3), padding='SAME')(x)
        x = nn.BatchNorm(use_running_average=not deterministic)(x)
        if residual.shape != x.shape:
            residual = nn.Conv(self.features, kernel_size=(1, 1))(residual)
        return nn.relu(x + residual)

class ResNet(nn.Module):
    num_classes: int = 10
    
    @nn.compact
    def __call__(self, x, deterministic=False):
        x = nn.Conv(64, kernel_size=(7, 7), strides=(2, 2), padding='SAME')(x)
        x = nn.BatchNorm(use_running_average=not deterministic)(x)
        x = nn.relu(x)
        x = nn.max_pool(x, window_shape=(3, 3), strides=(2, 2), padding='SAME')
        x = ResBlock(64)(x, deterministic=deterministic)
        x = ResBlock(128)(x, deterministic=deterministic)
        x = jnp.mean(x, axis=(1, 2))  # global average pooling
        x = nn.Dense(self.num_classes)(x)
        return x

# ============================================================
# 2. Training State (replaces model+optimizer combo)
# ============================================================
def create_train_state(key, model_cls, learning_rate=1e-3):
    # Initialize model variables
    dummy_input = jnp.ones((1, 32, 32, 3))
    variables = model_cls().init(key, dummy_input)
    params = variables['params']
    batch_stats = variables.get('batch_stats', {})
    
    # Setup optimizer
    tx = optax.chain(
        optax.clip_by_global_norm(1.0),
        optax.adamw(learning_rate, weight_decay=1e-4)
    )
    
    return {
        'params': params,
        'batch_stats': batch_stats,
        'opt_state': tx.init(params),
        'tx': tx,
        'step': 0,
        'key': key,
    }

# ============================================================
# 3. Loss and Metrics (pure functions)
# ============================================================
def compute_loss_and_accuracy(state, batch):
    images, labels = batch
    variables = {
        'params': state['params'], 
        'batch_stats': state['batch_stats']
    }
    
    # Forward pass (with mutable batch_stats for training)
    logits, updated_vars = ResNet(num_classes=10).apply(
        variables, images,
        mutable=['batch_stats'],
        deterministic=False  # training mode
    )
    
    loss = optax.softmax_cross_entropy_with_integer_labels(
        logits, labels).mean()
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == labels)
    
    return loss, accuracy, updated_vars['batch_stats']

# ============================================================
# 4. Training Step (JIT-compiled)
# ============================================================
@jax.jit
def train_step(state, batch):
    (loss, accuracy, new_batch_stats), grads = jax.value_and_grad(
        compute_loss_and_accuracy, argnums=0, has_aux=True
    )(state, batch)
    
    # Update parameters
    updates, new_opt_state = state['tx'].update(
        grads['params'], state['opt_state'], state['params']
    )
    new_params = optax.apply_updates(state['params'], updates)
    
    # Update state
    new_state = {
        'params': new_params,
        'batch_stats': new_batch_stats,
        'opt_state': new_opt_state,
        'tx': state['tx'],
        'step': state['step'] + 1,
        'key': state['key'],
    }
    
    return new_state, loss, accuracy

# ============================================================
# 5. Evaluation Step
# ============================================================
@jax.jit
def eval_step(state, batch):
    images, labels = batch
    variables = {
        'params': state['params'],
        'batch_stats': state['

🚀 Need a reliable AI agent for your project?

Deploy Hermes Agent in 10 minutes. Managed hosting, zero DevOps.

Get Started — $23.99/mo
← Back to all articles