Debugging Under JIT and Sharding

One of the most frustrating aspects of JAX is debugging JIT-compiled functions. Print statements inside @jax.jit only execute during tracing (once), not during actual computation. JAX provides special debugging tools to help.

jax.debug.print

import jax
import jax.numpy as jnp

@jax.jit
def buggy_function(x):
    # Regular print: only runs during tracing!
    print("This prints ONCE during tracing, not during execution")

    # jax.debug.print: runs during execution
    jax.debug.print("x = {x}", x=x)

    y = x ** 2
    jax.debug.print("y = {y}", y=y)

    # You can print inside conditions, scans, etc.
    z = jax.lax.cond(
        x.sum() > 0,
        lambda: (jax.debug.print("positive branch"), x * 2)[1],
        lambda: (jax.debug.print("negative branch"), x * -1)[1],
    )
    return z

result = buggy_function(jnp.array([1.0, -2.0, 3.0]))
# Prints:
# x = [1. -2. 3.]
# y = [1. 4. 9.]
# positive branch

jax.debug.breakpoint

@jax.jit
def function_with_breakpoint(x):
    y = jnp.sin(x)
    jax.debug.breakpoint()  # drops into pdb during execution
    z = jnp.cos(y)
    return z

# When this runs, you'll get an interactive pdb prompt
# where you can inspect y, x, etc.
# result = function_with_breakpoint(jnp.array(1.0))

Sharding and Multi-Device: Mesh, PartitionSpec, NamedSharding

For multi-device training, JAX uses a sharding system based on device meshes:

import jax
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding
from jax.experimental import mesh_utils
import jax.numpy as jnp

# Create a mesh of devices
devices = jax.devices()  # all available devices
mesh = Mesh(devices, axis_names=('data',))

# Or for multi-dimensional parallelism:
# devices = mesh_utils.create_device_mesh((2, 4))
# mesh = Mesh(devices, axis_names=('data', 'model'))

# Shard data across the 'data' axis
data_sharding = NamedSharding(mesh, P('data'))

# Replicate across all devices
replicated = NamedSharding(mesh, P())

# Note: jax.P is a convenient alias for PartitionSpec (since JAX 0.7.0)
data_sharding_v2 = NamedSharding(mesh, jax.P('data'))

# Place data on devices with sharding
x = jax.device_put(jnp.ones((1024, 256)), data_sharding)

# JIT with output sharding
@jax.jit
def forward(params, x):
    return x @ params

# Shard params: replicate across data axis
params = jax.device_put(jnp.ones((256, 128)), replicated)
output = forward(params, x)
print(output.sharding)  # shows how output is distributed

# Mixed precision is straightforward in JAX
@jax.jit
def train_step_bf16(params, x, y):
    def loss_fn(params):
        # Cast to bfloat16 for the forward pass
        x_bf16 = x.astype(jnp.bfloat16)
        logits = forward_bf16(params, x_bf16)
        return jnp.mean((logits.astype(jnp.float32) - y) ** 2)

    loss, grads = jax.value_and_grad(loss_fn)(params)
    # Gradients are in float32 — optimizer updates in full precision
    return loss, grads