Migrating from Legacy Frameworks to JAX

What is JAX?

JAX is a high-performance numerical computing library developed by Google that brings together automatic differentiation, hardware acceleration (GPU and TPU), and a NumPy-compatible API. Unlike legacy frameworks that bundle computation, graph building, and optimization into monolithic libraries, JAX is intentionally composable — it provides a set of small, sharp tools that can be combined in flexible ways.

At its core, JAX offers:

jax.numpy — a drop-in replacement for NumPy that runs on accelerators
Automatic differentiation — first-class grad, jacfwd, jacrev for arbitrary Python functions
JIT compilation — via jax.jit backed by XLA, Google's optimizing compiler
Vectorization — vmap for automatic batching without manual reshaping
Parallelization — pmap for multi-device data parallelism
Explicit random state — functional PRNG keys replace global random seeds

Legacy frameworks like TensorFlow 1.x, Theano, Caffe, or older Keras standalone code often entangle model definition, training loop, and hardware placement in ways that make debugging and customization difficult. JAX separates these concerns, giving you fine-grained control while still delivering production-grade performance.

Why Migrate to JAX?

🚀 Deploy your AI agent in 10 minutes

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

Try it free →

Migration isn't about chasing novelty — it's about solving concrete pain points that arise in legacy codebases:

Graph-building complexity: TensorFlow 1.x required explicit session management and graph construction. JAX uses eager execution everywhere, with optional JIT compilation that feels transparent.
Debugging friction: In static-graph frameworks, a shape mismatch error often surfaces far from its root cause. JAX executes Python functions directly, so standard debuggers and print statements work naturally.
Research-to-production gap: JAX's functional design means the same code that runs in a notebook can scale to TPU pods with minimal changes via pmap.
Deterministic randomness: Legacy global seeds cause reproducibility headaches across sessions and devices. JAX's explicit key-based PRNG eliminates this entire class of bugs.
Gradient computation flexibility: Need higher-order gradients, per-example gradients, or custom backward passes? JAX's grad composes with itself and with vmap, making these patterns trivial.

Migration Strategy: A Step-by-Step Guide

A successful migration proceeds incrementally. You don't need to rewrite an entire codebase at once. The following steps walk through the key transitions, from NumPy-heavy preprocessing to full training pipelines.

Step 1: Replace NumPy with jax.numpy

The gentlest entry point is swapping import numpy as np for import jax.numpy as jnp. Most array operations work identically, and you immediately gain the ability to run on GPU without changing any logic.

# Legacy NumPy code
import numpy as np

def preprocess(data):
    mean = np.mean(data, axis=0)
    centered = data - mean
    std = np.std(centered, axis=0)
    return centered / std

# JAX equivalent — same logic, now accelerator-ready
import jax.numpy as jnp

def preprocess(data):
    mean = jnp.mean(data, axis=0)
    centered = data - mean
    std = jnp.std(centered, axis=0)
    return centered / std

# To actually run on GPU, wrap with jax.jit
from jax import jit

@jit
def preprocess_fast(data):
    mean = jnp.mean(data, axis=0)
    centered = data - mean
    std = jnp.std(centered, axis=0)
    return centered / std

# Now call with a JAX array
key = jax.random.PRNGKey(42)
data = jax.random.normal(key, (1000, 128))
result = preprocess_fast(data)  # compiles and runs on GPU automatically

Key differences to watch for: JAX arrays are immutable — you cannot do arr[i] = value in-place. Use arr.at[i].set(value) for indexed updates, which returns a new array.

Step 2: Handle Random State Explicitly

Legacy frameworks often rely on a global random seed. In JAX, randomness is functional: you pass an explicit PRNG key and split it for each stochastic operation. This is initially verbose but eliminates non-determinism across parallel computations.

# Legacy approach (NumPy / TF1.x)
np.random.seed(42)
noise = np.random.normal(0, 1, (100,))

# JAX approach — explicit key management
from jax import random

key = random.PRNGKey(42)

# Split keys for different uses
key, subkey1, subkey2 = random.split(key, num=3)

noise = random.normal(subkey1, (100,))
dropout_mask = random.bernoulli(subkey2, 0.2, (100,))

# Key pattern inside a function that needs randomness
def stochastic_layer(params, x, key):
    key1, key2 = random.split(key)
    noise = random.normal(key1, x.shape)
    dropout = random.bernoulli(key2, 0.1, x.shape)
    return x + noise * dropout, key2  # return new key for downstream use

The pattern of returning a new key alongside your results is critical. Never reuse the same key — split before every stochastic operation. Libraries like flax and optax handle this internally when you pass a key to their high-level APIs.

Step 3: Convert Model Definitions to Pure Functions

Legacy Keras or PyTorch models use stateful layer objects. JAX encourages pure functions: parameters are passed in explicitly rather than stored inside the model. This makes parameter inspection, serialization, and gradient computation transparent.

# Legacy Keras model (stateful)
import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])
# Parameters live inside model — opaque to inspection

# JAX + Flax equivalent (functional, explicit params)
import flax.linen as nn

class MLP(nn.Module):
    hidden_size: int = 128
    output_size: int = 10

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(self.hidden_size)(x)
        x = nn.relu(x)
        x = nn.Dense(self.output_size)(x)
        return x

# Initialize parameters explicitly
model = MLP()
key = jax.random.PRNGKey(0)
x_dummy = jnp.ones((1, 64))
params = model.init(key, x_dummy)

# Forward pass — params passed explicitly
def forward(params, x):
    return model.apply(params, x)

output = forward(params, x_dummy)
# params is a plain nested dict; you can inspect, save, or share it freely

This separation means you can write loss functions that take parameters as their first argument — exactly what jax.grad expects.

Step 4: Replace Legacy Optimizers with Optax

Optax is the standard optimization library for JAX. It provides composable gradient transformations: you chain clipping, weight decay, and learning rate schedules together, then apply the resulting optimizer as a pure function.

# Legacy optimizer (TensorFlow 1.x)
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
train_op = optimizer.minimize(loss)

# Optax equivalent — optimizer is a gradient transform
import optax

optimizer = optax.adam(learning_rate=0.001)

# Compose with gradient clipping and weight decay
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=0.001, weight_decay=1e-4),
)

# Usage inside a training step
opt_state = optimizer.init(params)

def train_step(params, opt_state, x, y):
    loss_value, 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_value

Optax's functional design means you can inspect or modify gradients between computation and application. You can even write custom gradient transforms that integrate seamlessly.

Step 5: Build Training Loops with JIT and vmap

Here's where JAX shines. A complete training loop combines jit for compilation, vmap for batching, and grad for differentiation — all as orthogonal transformations of pure functions.

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

# Model definition
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 = x.reshape((x.shape[0], -1))
        x = nn.Dense(10)(x)
        return x

# Loss function (pure, takes params as first arg)
def loss_fn(params, x, y):
    logits = model.apply(params, x)
    loss = optax.softmax_cross_entropy_with_integer_labels(logits, y)
    return loss.mean()

# Single training step — will be JIT-compiled
@jax.jit
def train_step(params, opt_state, x_batch, y_batch):
    loss, grads = jax.value_and_grad(loss_fn)(params, x_batch, y_batch)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

# Evaluation step — uses vmap for automatic batching
@jax.jit
def evaluate(params, x, y):
    # vmap the loss computation over the batch dimension
    batch_losses = jax.vmap(
        lambda x_single, y_single: loss_fn(params, x_single[None], y_single[None])
    )(x, y)
    return batch_losses.mean()

# Full training loop
model = CNN()
key = jax.random.PRNGKey(0)
dummy_x = jnp.ones((1, 28, 28, 1))
params = model.init(key, dummy_x)

optimizer = optax.adamw(learning_rate=0.001, weight_decay=1e-4)
opt_state = optimizer.init(params)

# Training
for epoch in range(10):
    for batch_x, batch_y in dataloader:
        params, opt_state, loss = train_step(params, opt_state, batch_x, batch_y)
    # Evaluate after each epoch
    eval_loss = evaluate(params, x_test, y_test)
    print(f"Epoch {epoch}: train_loss={loss:.4f}, eval_loss={eval_loss:.4f}")

Step 6: Scale with pmap for Multi-Device Parallelism

When moving from a single GPU to multiple devices (GPU pods or TPU slices), JAX's pmap transforms your single-device function into a parallel version with minimal code changes. This is dramatically simpler than legacy distributed training setups.

# Single-device function
@jax.jit
def single_device_step(params, x, y):
    grads = jax.grad(loss_fn)(params, x, y)
    return grads

# Multi-device version with pmap
from jax import pmap

# Replicate parameters across devices
devices = jax.devices()  # e.g., 8 GPU cores
params_replicated = jax.device_put_replicated(params, devices)

@pmap
def multi_device_step(params, x, y):
    # Each device gets a slice of x, y
    grads = jax.grad(loss_fn)(params, x, y)
    # Average gradients across devices
    grads = jax.lax.pmean(grads, axis_name='devices')
    return grads

# Data is automatically sharded across devices by pmap
# Just reshape your batch to (num_devices, batch_per_device, ...)
batch_reshaped = batch.reshape(jax.device_count(), -1, *batch.shape[1:])
grads = multi_device_step(params_replicated, batch_reshaped, labels_reshaped)

Common Migration Pitfalls and Solutions

Pitfall 1: In-Place Mutations

Legacy code often mutates arrays in-place. JAX arrays are immutable. The fix: use .at[...] syntax or rewrite operations to produce new arrays.

# Won't work in JAX
# arr[0, 1] = 5.0

# Correct JAX approach
arr = arr.at[0, 1].set(5.0)

# For conditional updates
arr = arr.at[arr > 0].set(0.0)  # clip positive values to zero

Pitfall 2: Python Side Effects Inside JIT

jax.jit traces your function and compiles it. Any Python-only operations (print, if-statements on non-array values, external state mutations) happen at trace time, not at runtime. Use jax.lax.cond, jax.lax.scan, etc., for control flow that depends on array values.

# Problematic: Python if on array value
@jax.jit
def problematic(x):
    if x.sum() > 0:  # This is evaluated once at trace time!
        return x * 2
    return x

# Correct: use jax.lax.cond for array-dependent branching
@jax.jit
def correct(x):
    return jax.lax.cond(
        x.sum() > 0,
        lambda: x * 2,
        lambda: x,
    )

Pitfall 3: Forgetting to Split PRNG Keys

Reusing the same key across two random operations produces identical "random" outputs. Always split.

# Dangerous — dropout1 and dropout2 will be identical
key = jax.random.PRNGKey(42)
dropout1 = jax.random.bernoulli(key, 0.5, shape)
dropout2 = jax.random.bernoulli(key, 0.5, shape)  # same key!

# Safe pattern
key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)
dropout1 = jax.random.bernoulli(key1, 0.5, shape)
dropout2 = jax.random.bernoulli(key2, 0.5, shape)  # different, independent

Pitfall 4: Large JIT Compilation Overheads on First Run

JIT compilation happens on the first call with a given input shape. Subsequent calls with the same shape are fast. If your data shapes change frequently (e.g., variable-length sequences), consider padding to fixed sizes or using jax.lax.scan for dynamic-length processing.

Best Practices for Long-Term JAX Development

Keep functions pure: A function should take all its inputs as arguments and return all outputs. No global state, no side effects. This makes jit, grad, and vmap always applicable.
Structure parameters as nested dictionaries: Flax and other JAX libraries use dict-like parameter trees. Serialize them with flax.serialization or plain pickle. They're just Python objects.
Use jax.tree_util for parameter manipulation: Functions like jax.tree_map let you apply operations across nested parameter structures without manual recursion.
Profile before scaling: Use jax.profiler and the --logtostderr XLA flags to understand compilation times and memory usage before adding more devices.
Adopt Flax or Haiku for complex models: Writing raw parameter management works for small models but becomes tedious. Flax provides nn.Module with init/apply while staying fully functional.
Test gradient correctness: Use jax.test_utils.check_grads to verify your custom gradients against finite differences, especially after migrating complex loss functions.
Version your parameter checkpoints: Since parameters are plain dicts, store them alongside a version tag and the model code hash. This prevents silent mismatches when model architecture changes.

Complete Migration Example: From TensorFlow 1.x to JAX

Below is a side-by-side comparison of a full MNIST training pipeline — first in TensorFlow 1.x style, then fully migrated to JAX with Flax and Optax.

# =============================================
# LEGACY: TensorFlow 1.x MNIST Classifier
# =============================================
import tensorflow as tf

# Graph construction
tf.reset_default_graph()
X = tf.placeholder(tf.float32, [None, 784])
y = tf.placeholder(tf.int32, [None])

W1 = tf.Variable(tf.random_normal([784, 128]))
b1 = tf.Variable(tf.zeros([128]))
h = tf.nn.relu(tf.matmul(X, W1) + b1)

W2 = tf.Variable(tf.random_normal([128, 10]))
b2 = tf.Variable(tf.zeros([10]))
logits = tf.matmul(h, W2) + b2

loss = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=logits)
)
train_op = tf.train.AdamOptimizer(0.001).minimize(loss)

# Session management
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for epoch in range(10):
        for batch_x, batch_y in train_data:
            _, l = sess.run(
                [train_op, loss],
                feed_dict={X: batch_x, y: batch_y}
            )
        print(f"Epoch {epoch}, loss: {l}")


# =============================================
# MIGRATED: JAX + Flax + Optax MNIST Classifier
# =============================================
import jax
import jax.numpy as jnp
import optax
from flax import linen as nn
from flax.training import train_state

# Pure model definition
class MNISTClassifier(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(128)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

# Pure loss function
def compute_loss(params, x, y):
    logits = model.apply(params, x)
    loss = optax.softmax_cross_entropy_with_integer_labels(logits, y)
    return loss.mean()

# Single training step — pure, JIT-friendly
@jax.jit
def train_step(state, x_batch, y_batch):
    def loss_fn(params):
        return compute_loss(params, x_batch, y_batch)
    loss, grads = jax.value_and_grad(loss_fn)(state.params)
    state = state.apply_gradients(grads=grads)
    return state, loss

# Evaluation
@jax.jit
def evaluate(params, x, y):
    logits = model.apply(params, x)
    predictions = jnp.argmax(logits, axis=-1)
    accuracy = jnp.mean(predictions == y)
    return accuracy

# Initialization
model = MNISTClassifier()
key = jax.random.PRNGKey(0)
dummy_x = jnp.ones((1, 784))
params = model.init(key, dummy_x)

optimizer = optax.adam(learning_rate=0.001)
state = train_state.TrainState.create(
    apply_fn=model.apply,
    params=params,
    tx=optimizer,
)

# Training — no sessions, no placeholders
for epoch in range(10):
    for batch_x, batch_y in train_data:
        state, loss = train_step(state, batch_x, batch_y)
    acc = evaluate(state.params, test_x, test_y)
    print(f"Epoch {epoch}: loss={loss:.4f}, accuracy={acc:.4f}")

Conclusion

Migrating from legacy frameworks to JAX is a process of progressively disentangling computation from state, side effects from pure logic, and hardware placement from algorithm design. The initial investment — explicit parameter passing, functional randomness, immutable arrays — pays off in dramatically simpler debugging, transparent scaling across devices, and composability that legacy graph-based frameworks simply cannot match.

Start small: replace NumPy with jax.numpy in data preprocessing, then convert one model at a time using Flax, and finally wrap your training loop with jit and vmap. Each step brings immediate benefits in performance and code clarity, and the full migration leaves you with a codebase where every component — model, optimizer, loss, data pipeline — is a testable, reusable pure function. That functional foundation makes your research code production-ready by default, closing the gap that legacy frameworks forced you to bridge with ad hoc engineering.