Testing JAX Applications: Unit Tests to Integration

Testing JAX Applications: Unit Tests to Integration

JAX has rapidly become one of the most powerful libraries for high-performance numerical computing and machine learning research. Its functional programming model, automatic differentiation, and hardware acceleration make it ideal for building everything from simple numerical solvers to complex neural networks. However, JAX's unique design—particularly its functional nature and the use of transformations like jit, vmap, and grad—introduces specific testing challenges that traditional Python testing patterns don't always address well.

In this tutorial, you'll learn how to build a robust testing strategy for JAX applications, starting from pure function unit tests and progressing all the way to integration tests that validate end-to-end training pipelines. We'll cover the tools, patterns, and best practices that will help you catch bugs early, ensure numerical correctness, and maintain confidence as your codebase grows.

Why Testing JAX Applications Matters

JAX code often involves subtle numerical operations where small mistakes can lead to silently incorrect results. A transposed matrix, a wrong axis in a reduction, or an off-by-one in shape manipulation can produce outputs that look plausible but are mathematically wrong. Unlike typical web applications where errors often manifest as exceptions, numerical bugs frequently produce valid-looking arrays with incorrect values.

• Numerical correctness: Gradients, matrix operations, and reductions must produce mathematically correct results within acceptable tolerances.
• Shape safety: JAX's shape system catches many errors at trace time, but some only surface during specific execution paths.
• Transformation correctness: Composing jit, grad, vmap, and pmap can introduce subtle bugs that only appear under specific conditions.
• Reproducibility: Random number generation in JAX is stateless, which is great for reproducibility but requires explicit key management in tests.
• Performance regression: Unintended recompilation or host-device transfers can silently degrade performance.

Setting Up Your Testing Environment

Before diving into test patterns, let's establish a solid testing setup. The standard tools for JAX testing are pytest for test orchestration and jax.test_util (where available) along with NumPy's testing utilities for numerical assertions.

# requirements-test.txt
pytest>=7.0
pytest-xdist>=3.0
jax>=0.4.0
jaxlib>=0.4.0
numpy>=1.20

Create a conftest.py file to configure JAX for testing. This ensures consistent behavior across your test suite and can disable JIT for easier debugging when needed.

# conftest.py
import os
import pytest
import jax

# Enable 64-bit precision for tests that need it
jax.config.update("jax_enable_x64", True)

# Optionally disable JIT for easier debugging
# os.environ["JAX_DISABLE_JIT"] = "True"


@pytest.fixture(autouse=True)
def reset_rng_state():
    """Reset any global state between tests."""
    yield


def pytest_configure(config):
    """Configure JAX for testing."""
    os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

Unit Testing Pure JAX Functions

The foundation of JAX testing is unit testing pure functions. Because JAX functions are pure (no side effects), they are inherently testable—you provide inputs, check outputs, and don't need to worry about internal state. Let's start with a simple example.

# src/layers.py
import jax
import jax.numpy as jnp


def linear(x, w, b):
    """A simple linear layer: y = x @ w + b."""
    return jnp.dot(x, w) + b


def relu(x):
    """ReLU activation function."""
    return jnp.maximum(x, 0.0)


def softmax(x, axis=-1):
    """Numerically stable softmax."""
    x_max = jnp.max(x, axis=axis, keepdims=True)
    exp_x = jnp.exp(x - x_max)
    return exp_x / jnp.sum(exp_x, axis=axis, keepdims=True)

Now let's write unit tests for these functions. The key principle is to test both correctness (are the values right?) and properties (does the function satisfy mathematical invariants?).

# tests/test_layers.py
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from src.layers import linear, relu, softmax


class TestLinear:
    def test_basic_multiplication(self):
        x = jnp.ones((2, 3))
        w = jnp.ones((3, 4))
        b = jnp.zeros((4,))
        result = linear(x, w, b)
        assert result.shape == (2, 4)
        np.testing.assert_allclose(result, jnp.full((2, 4), 3.0))

    def test_bias_addition(self):
        x = jnp.zeros((1, 2))
        w = jnp.ones((2, 2))
        b = jnp.array([1.0, 2.0])
        result = linear(x, w, b)
        np.testing.assert_allclose(result, jnp.array([[1.0, 2.0]]))

    def test_with_random_inputs(self):
        key = jax.random.PRNGKey(42)
        k1, k2, k3 = jax.random.split(key, 3)
        x = jax.random.normal(k1, (10, 5))
        w = jax.random.normal(k2, (5, 3))
        b = jax.random.normal(k3, (3,))
        result = linear(x, w, b)
        assert result.shape == (10, 3)

    def test_gradient_computation(self):
        """Verify gradients are computed correctly."""
        x = jnp.array([[1.0, 2.0], [3.0, 4.0]])
        w = jnp.array([[1.0], [1.0]])
        b = jnp.array([0.0])

        def loss_fn(w, b):
            return jnp.sum(linear(x, w, b) ** 2)

        grad_w, grad_b = jax.grad(loss_fn, argnums=(0, 1))(w, b)
        # d/dw of sum((x@w)^2) = 2 * x.T @ (x @ w)
        expected_grad_w = 2 * x.T @ (x @ w)
        np.testing.assert_allclose(grad_w, expected_grad_w, rtol=1e-5)


class TestReLU:
    def test_positive_values(self):
        x = jnp.array([1.0, 2.0, 3.0])
        np.testing.assert_allclose(relu(x), x)

    def test_negative_values(self):
        x = jnp.array([-1.0, -2.0, -3.0])
        np.testing.assert_allclose(relu(x), jnp.zeros_like(x))

    def test_mixed_values(self):
        x = jnp.array([-1.0, 0.0, 1.0])
        np.testing.assert_allclose(relu(x), jnp.array([0.0, 0.0, 1.0]))

    def test_zero_gradient(self):
        """Gradient at zero should be 0 in JAX's convention."""
        grad_fn = jax.grad(lambda x: relu(x).sum())
        assert grad_fn(jnp.array(-1.0)) == 0.0
        assert grad_fn(jnp.array(1.0)) == 1.0


class TestSoftmax:
    def test_outputs_sum_to_one(self):
        key = jax.random.PRNGKey(0)
        x = jax.random.normal(key, (5, 10))
        result = softmax(x)
        sums = jnp.sum(result, axis=-1)
        np.testing.assert_allclose(sums, jnp.ones((5,)), atol=1e-6)

    def test_all_outputs_positive(self):
        x = jnp.array([[-100.0, -200.0, -300.0]])
        result = softmax(x)
        assert jnp.all(result > 0)

    def test_numerical_stability(self):
        """Softmax should handle very large values without overflow."""
        x = jnp.array([1000.0, 1001.0, 1002.0])
        result = softmax(x)
        assert jnp.all(jnp.isfinite(result))
        np.testing.assert_allclose(jnp.sum(result), 1.0, atol=1e-6)

    def test_uniform_input(self):
        """Uniform input should produce uniform output."""
        x = jnp.array([5.0, 5.0, 5.0])
        result = softmax(x)
        np.testing.assert_allclose(result, jnp.ones(3) / 3.0)

Testing JAX Transformations

JAX's power comes from its transformations—jit, grad, vmap, and pmap. Each transformation has its own testing considerations. Let's explore how to test them effectively.

# src/transforms.py
import jax
import jax.numpy as jnp


@jax.jit
def batched_distance(x, y):
    """Compute pairwise distances between two sets of vectors."""
    diff = x[:, None, :] - y[None, :, :]
    return jnp.sqrt(jnp.sum(diff ** 2, axis=-1))


def gradient_norm(params, loss_fn):
    """Compute the L2 norm of the gradient."""
    grads = jax.grad(loss_fn)(params)
    flat_grads, _ = jax.flatten_util.ravel_pytree(grads)
    return jnp.linalg.norm(flat_grads)


def hessian_diagonal(fn, x):
    """Compute the diagonal of the Hessian of fn at x."""
    grad_fn = jax.grad(fn)
    def single_diag(xi, i):
        return jax.grad(lambda x: grad_fn(x)[i])(x)[i]
    return jax.vmap(single_diag, in_axes=(None, 0))(x, jnp.arange(x.shape[0]))

# tests/test_transforms.py
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from src.transforms import batched_distance, gradient_norm, hessian_diagonal


class TestJIT:
    def test_jit_preserves_results(self):
        key = jax.random.PRNGKey(0)
        x = jax.random.normal(key, (5, 3))
        y = jax.random.normal(key, (4, 3))

        # Non-jitted version
        def distance_fn(x, y):
            diff = x[:, None, :] - y[None, :, :]
            return jnp.sqrt(jnp.sum(diff ** 2, axis=-1))

        result_jit = batched_distance(x, y)
        result_no_jit = distance_fn(x, y)
        np.testing.assert_allclose(result_jit, result_no_jit, rtol=1e-6)

    def test_jit_with_different_shapes(self):
        """JIT should handle different input shapes via recompilation."""
        x1 = jnp.ones((3, 2))
        y1 = jnp.ones((2, 2))
        r1 = batched_distance(x1, y1)
        assert r1.shape == (3, 2)

    def test_jit_static_argnames(self):
        """Test functions with static arguments."""
        @jax.jit
        def power_fn(x, n):
            return x ** n

        # This works because n is traced as a static value
        result = jax.jit(lambda x: x ** 3)(jnp.array(2.0))
        assert result == 8.0


class TestGrad:
    def test_gradient_correctness(self):
        """Test gradient of a known function."""
        def fn(x):
            return x ** 2 + 3 * x + 1

        x = jnp.array(2.0)
        grad = jax.grad(fn)(x)
        # d/dx (x^2 + 3x + 1) = 2x + 3
        np.testing.assert_allclose(grad, 7.0)

    def test_gradient_of_constant_is_zero(self):
        grad = jax.grad(lambda x: jnp.array(5.0))(jnp.array(3.0))
        np.testing.assert_allclose(grad, 0.0)

    def test_higher_order_gradients(self):
        def fn(x):
            return x ** 4

        x = jnp.array(2.0)
        first = jax.grad(fn)(x)       # 4x^3 = 32
        second = jax.grad(jax.grad(fn))(x)  # 12x^2 = 48
        third = jax.grad(jax.grad(jax.grad(fn)))(x)  # 24x = 48

        np.testing.assert_allclose(first, 32.0)
        np.testing.assert_allclose(second, 48.0)
        np.testing.assert_allclose(third, 48.0)

    def test_gradient_norm(self):
        params = {"w": jnp.array([3.0, 4.0])}
        loss_fn = lambda p: jnp.sum(p["w"] ** 2)
        norm = gradient_norm(params, loss_fn)
        # grad = [6, 8], norm = 10
        np.testing.assert_allclose(norm, 10.0, rtol=1e-5)


class TestVmap:
    def test_vmap_matches_loop(self):
        def fn(x):
            return x ** 2 + 1

        x = jnp.arange(5.0)
        vmapped = jax.vmap(fn)(x)
        looped = jnp.array([fn(xi) for xi in x])
        np.testing.assert_allclose(vmapped, looped)

    def test_vmap_with_axes(self):
        def dot(a, b):
            return jnp.dot(a, b)

        a = jnp.ones((3, 4))
        b = jnp.ones((4,))

        # Map over first axis of a
        result = jax.vmap(dot, in_axes=(0, None))(a, b)
        np.testing.assert_allclose(result, jnp.full((3,), 4.0))

    def test_hessian_diagonal(self):
        def fn(x):
            return jnp.sum(x ** 4)

        x = jnp.array([1.0, 2.0, 3.0])
        diag = hessian_diagonal(fn, x)
        # d^2/dx_i^2 of sum(x_i^4) = 12 * x_i^2
        expected = 12 * x ** 2
        np.testing.assert_allclose(diag, expected, rtol=1e-5)

Testing with PyTrees

JAX uses pytrees extensively for managing parameters, optimizer states, and nested data structures. Testing code that manipulates pytrees requires understanding how to construct, traverse, and compare them.

# src/model.py
import jax
import jax.numpy as jnp
from typing import NamedTuple


class MLPParams(NamedTuple):
    w1: jnp.ndarray
    b1: jnp.ndarray
    w2: jnp.ndarray
    b2: jnp.ndarray


def init_params(key, input_dim, hidden_dim, output_dim):
    k1, k2 = jax.random.split(key)
    w1 = jax.random.normal(k1, (input_dim, hidden_dim)) * 0.1
    b1 = jnp.zeros(hidden_dim)
    w2 = jax.random.normal(k2, (hidden_dim, output_dim)) * 0.1
    b2 = jnp.zeros(output_dim)
    return MLPParams(w1, b1, w2, b2)


def mlp_forward(params, x):
    h = jnp.tanh(jnp.dot(x, params.w1) + params.b1)
    return jnp.dot(h, params.w2) + params.b2


def cross_entropy_loss(params, x, y, num_classes):
    logits = mlp_forward(params, x)
    log_probs = jax.nn.log_softmax(logits)
    one_hot = jax.nn.one_hot(y, num_classes)
    return -jnp.mean(jnp.sum(one_hot * log_probs, axis=-1))

# tests/test_model.py
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from src.model import init_params, mlp_forward, cross_entropy_loss, MLPParams


class TestPytreeParams:
    def test_params_are_pytree(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 5)
        leaves, treedef = jax.tree_util.tree_flatten(params)
        assert len(leaves) == 4  # w1, b1, w2, b2

    def test_params_roundtrip(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 5)
        leaves, treedef = jax.tree_util.tree_flatten(params)
        reconstructed = jax.tree_util.tree_unflatten(treedef, leaves)
        for a, b in zip(params, reconstructed):
            np.testing.assert_array_equal(a, b)

    def test_grad_returns_same_structure(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 5)
        x = jax.random.normal(key, (4, 10))
        y = jnp.array([0, 1, 2, 3])

        loss_fn = lambda p: cross_entropy_loss(p, x, y, 5)
        grads = jax.grad(loss_fn)(params)

        # Gradients should have the same structure as params
        grad_leaves, grad_treedef = jax.tree_util.tree_flatten(grads)
        _, param_treedef = jax.tree_util.tree_flatten(params)
        assert grad_treedef == param_treedef
        assert len(grad_leaves) == 4

    def test_tree_map_operations(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 5)

        # Scale all parameters
        scaled = jax.tree_util.tree_map(lambda x: x * 2.0, params)
        for p, s in zip(params, scaled):
            np.testing.assert_allclose(s, p * 2.0)

    def test_tree_map_with_two_inputs(self):
        key = jax.random.PRNGKey(0)
        params1 = init_params(key, 10, 20, 5)
        params2 = init_params(jax.random.PRNGKey(1), 10, 20, 5)

        diff = jax.tree_util.tree_map(lambda a, b: a - b, params1, params2)
        for d, a, b in zip(diff, params1, params2):
            np.testing.assert_allclose(d, a - b)


class TestMLPForward:
    def test_output_shape(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 5)
        x = jax.random.normal(key, (8, 10))
        output = mlp_forward(params, x)
        assert output.shape == (8, 5)

    def test_batched_and_single_match(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 5)
        x = jax.random.normal(key, (4, 10))

        batched = mlp_forward(params, x)
        single = jax.vmap(lambda xi: mlp_forward(params, xi))(x)
        np.testing.assert_allclose(batched, single, rtol=1e-5)


class TestCrossEntropy:
    def test_loss_is_non_negative(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 3)
        x = jax.random.normal(key, (8, 10))
        y = jax.random.randint(key, (8,), 0, 3)
        loss = cross_entropy_loss(params, x, y, 3)
        assert loss >= 0

    def test_perfect_prediction_low_loss(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 4, 8, 3)
        x = jnp.eye(4)
        y = jnp.array([0, 1, 2, 0])

        # Train briefly to get low loss
        loss_fn = lambda p: cross_entropy_loss(p, x, y, 3)
        grads = jax.grad(loss_fn)(params)
        for _ in range(100):
            params = jax.tree_util.tree_map(
                lambda p, g: p - 0.5 * g, params, grads
            )
            grads = jax.grad(loss_fn)(params)

        loss = cross_entropy_loss(params, x, y, 3)
        assert loss < 0.1

Property-Based Testing with JAX

Property-based testing is especially powerful for numerical code because it can automatically generate test cases and find edge cases you might miss. Using hypothesis with JAX, you can test mathematical invariants that should hold for any valid input.

# tests/test_properties.py
import jax
import jax.numpy as jnp
import numpy as np
from hypothesis import given, strategies as st, settings
from hypothesis.extra.numpy import arrays
from src.layers import softmax, relu


@given(
    arrays(dtype=np.float32, shape=(5, 10),
           elements=st.floats(-10, 10, allow_nan=False, allow_infinity=False))
)
@settings(max_examples=50)
def test_softmax_always_sums_to_one(x):
    result = softmax(jnp.array(x))
    sums = jnp.sum(result, axis=-1)
    np.testing.assert_allclose(sums, np.ones(5), atol=1e-5)


@given(
    arrays(dtype=np.float32, shape=(100,),
           elements=st.floats(-1000, 1000, allow_nan=False, allow_infinity=False))
)
@settings(max_examples=50)
def test_relu_never_negative(x):
    result = relu(jnp.array(x))
    assert jnp.all(result >= 0)


@given(
    arrays(dtype=np.float32, shape=(3, 3),
           elements=st.floats(-5, 5, allow_nan=False, allow_infinity=False))
)
@settings(max_examples=50)
def test_gradient_of_sum_is_one(x):
    """The gradient of sum(x) with respect to x should be all ones."""
    x_jax = jnp.array(x)
    grad = jax.grad(lambda v: jnp.sum(v))(x_jax)
    np.testing.assert_allclose(grad, np.ones((3, 3)), atol=1e-5)


@given(
    st.integers(1, 50),
    st.integers(1, 50),
)
@settings(max_examples=20)
def test_matmul_shape_rule(a, b):
    """Matrix multiplication shape rule: (a, b) @ (b, c) -> (a, c)."""
    c = 10
    key = jax.random.PRNGKey(0)
    x = jax.random.normal(key, (a, b))
    w = jax.random.normal(key, (b, c))
    result = jnp.dot(x, w)
    assert result.shape == (a, c)

Testing Optimizers and Training Steps

Optimizers are critical components that need careful testing. You should verify that they correctly update parameters, maintain state, and actually reduce loss over time.

# src/optimizer.py
import jax
import jax.numpy as jnp
from typing import NamedTuple


class AdamState(NamedTuple):
    m: any  # first moment
    v: any  # second moment
    t: int  # timestep


def adam_init(params):
    return AdamState(
        m=jax.tree_util.tree_map(jnp.zeros_like, params),
        v=jax.tree_util.tree_map(jnp.zeros_like, params),
        t=0,
    )


def adam_step(params, grads, state, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8):
    t = state.t + 1
    m = jax.tree_util.tree_map(
        lambda m_i, g_i: beta1 * m_i + (1 - beta1) * g_i, state.m, grads
    )
    v = jax.tree_util.tree_map(
        lambda v_i, g_i: beta2 * v_i + (1 - beta2) * g_i ** 2, state.v, grads
    )
    m_hat = jax.tree_util.tree_map(lambda m_i: m_i / (1 - beta1 ** t), m)
    v_hat = jax.tree_util.tree_map(lambda v_i: v_i / (1 - beta2 ** t), v)
    new_params = jax.tree_util.tree_map(
        lambda p, m_i, v_i: p - lr * m_i / (jnp.sqrt(v_i) + eps),
        params, m_hat, v_hat
    )
    return new_params, AdamState(m, v, t)

# tests/test_optimizer.py
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from src.optimizer import adam_init, adam_step


class TestAdam:
    def test_state_initialization(self):
        params = {"w": jnp.ones((3, 3)), "b": jnp.zeros(3)}
        state = adam_init(params)
        assert state.t == 0
        np.testing.assert_array_equal(state.m["w"], jnp.zeros((3, 3)))
        np.testing.assert_array_equal(state.v["b"], jnp.zeros(3))

    def test_parameters_change(self):
        params = {"w": jnp.array([1.0, 2.0, 3.0])}
        grads = {"w": jnp.array([0.1, 0.1, 0.1])}
        state = adam_init(params)
        new_params, new_state = adam_step(params, grads, state)
        assert new_state.t == 1
        assert not np.allclose(new_params["w"], params["w"])

    def test_zero_gradient_no_change(self):
        params = {"w": jnp.array([1.0, 2.0, 3.0])}
        grads = {"w": jnp.array([0.0, 0.0, 0.0])}
        state = adam_init(params)
        new_params, _ = adam_step(params, grads, state)
        np.testing.assert_allclose(new_params["w"], params["w"])

    def test_loss_decreases_on_quadratic(self):
        """Adam should minimize a simple quadratic function."""
        def loss_fn(w):
            return jnp.sum((w - jnp.array([5.0, 5.0, 5.0])) ** 2)

        params = {"w": jnp.array([0.0, 0.0, 0.0])}
        state = adam_init(params)

        initial_loss = loss_fn(params["w"])
        for _ in range(500):
            grads = jax.grad(loss_fn)(params["w"])
            grads_dict = {"w": grads}
            params, state = adam_step(params, grads_dict, state, lr=0.1)

        final_loss = loss_fn(params["w"])
        assert final_loss < initial_loss
        assert final_loss < 0.01

    def test_jitted_training_step(self):
        """The optimizer step should be jittable."""
        params = {"w": jnp.ones((2, 2))}
        grads = {"w": jnp.ones((2, 2)) * 0.1}
        state = adam_init(params)

        jitted_step = jax.jit(adam_step)
        new_params, new_state = jitted_step(params, grads, state)
        assert new_state.t == 1
        assert new_params["w"].shape == (2, 2)

Integration Testing: End-to-End Training Pipelines

Integration tests verify that multiple components work together correctly. For JAX applications, this typically means testing the full training loop: data loading, forward pass, loss computation, gradient computation, parameter updates, and evaluation.

# src/training.py
import jax
import jax.numpy as jnp
from src.model import init_params, mlp_forward, cross_entropy_loss
from src.optimizer import adam_init, adam_step


def compute_accuracy(params, x, y, num_classes):
    logits = mlp_forward(params, x)
    predictions = jnp.argmax(logits, axis=-1)
    return jnp.mean(predictions == y)


def train_step(params, state, x, y, num_classes, lr):
    def loss_fn(p):
        return cross_entropy_loss(p, x, y, num_classes)
    grads = jax.grad(loss_fn)(params)
    new_params, new_state = adam_step(params, grads, state, lr=lr)
    loss = loss_fn(new_params)
    return new_params, new_state, loss


def train_epoch(params, state, dataset, num_classes, lr):
    """Train for one epoch over the dataset."""
    jit_step = jax.jit(train_step)
    total_loss = 0.0
    num_batches = 0
    for x, y in dataset:
        params, state, loss = jit_step(params, state, x, y, num_classes, lr)
        total_loss += float(loss)
        num_batches += 1
    return params, state, total_loss / max(num_batches, 1)


def train_model(params, train_data, val_data, num_classes,
                epochs=10, lr=1e-3):
    """Full training loop with validation."""
    state = adam_init(params)
    history = {"train_loss": [], "val_loss": [], "val_acc": []}

    for epoch in range(epochs):
        params, state, avg_loss = train_epoch(
            params, state, train_data, num_classes, lr
        )
        val_x, val_y = val_data
        val_loss = float(cross_entropy_loss(params, val_x, val_y, num_classes))
        val_acc = float(compute_accuracy(params, val_x, val_y, num_classes))

        history["train_loss"].append(avg_loss)
        history["val_loss"].append(val_loss)
        history["val_acc"].append(val_acc)

    return params, state, history

# tests/test_training_integration.py
import jax
import jax.numpy as jnp
import numpy as np
import pytest
from src.model import init_params, mlp_forward, cross_entropy_loss
from src.training import train_step, train_epoch, train_model, compute_accuracy
from src.optimizer import adam_init


def make_synthetic_dataset(key, num_samples, input_dim, num_classes, batch_size):
    """Generate a linearly separable synthetic dataset."""
    k1, k2 = jax.random.split(key)
    X = jax.random.normal(k1, (num_samples, input_dim))
    # Make labels based on a linear separator
    true_w = jax.random.normal(k2, (input_dim, num_classes))
    logits = X @ true_w
    y = jnp.argmax(logits, axis=-1)

    batches = []
    for i in range(0, num_samples, batch_size):
        batches.append((X[i:i+batch_size], y[i:i+batch_size]))
    return batches


class TestTrainStep:
    def test_returns_correct_types(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 3)
        state = adam_init(params)
        x = jax.random.normal(key, (8, 10))
        y = jax.random.randint(key, (8,), 0, 3)

        new_params, new_state, loss = train_step(
            params, state, x, y, 3, lr=1e-3
        )
        assert isinstance(loss, jnp.ndarray) or isinstance(loss, float)
        assert new_state.t == state.t + 1

    def test_loss_decreases_with_multiple_steps(self):
        key = jax.random.PRNGKey(42)
        params = init_params(key, 10, 20, 3)
        state = adam_init(params)

        dataset = make_synthetic_dataset(key, 64, 10, 3, 16)
        x, y = dataset[0]

        initial_loss = float(cross_entropy_loss(params, x, y, 3))
        for _ in range(50):
            params, state, loss = train_step(params, state, x, y, 3, lr=1e-2)
        final_loss = float(loss)

        assert final_loss < initial_loss


class TestTrainEpoch:
    def test_processes_all_batches(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 10, 20, 3)
        state = adam_init(params)
        dataset = make_synthetic_dataset(key, 64, 10, 3, 16)

        params, state, avg_loss = train_epoch(
            params, state, dataset, 3, lr=1e-3
        )
        assert avg_loss > 0
        assert isinstance(avg_loss, float)


class TestTrainModel:
    @pytest.mark.slow
    def test_full_training_improves_accuracy(self):
        """End-to-end test: model should learn to classify synthetic data."""
        key = jax.random.PRNGKey(123)
        params = init_params(key, 10, 32, 3)

        train_data = make_synthetic_dataset(key, 128, 10, 3, 32)
        val_x = jax.random.normal(jax.random.PRNGKey(456), (32, 10))
        val_y = jax.random.randint(jax.random.PRNGKey(789), (32,), 0, 3)
        val_data = (val_x, val_y)

        initial_acc = float(compute_accuracy(params, val_x, val_y, 3))

        params, state, history = train_model(
            params, train_data, val_data, num_classes=3,
            epochs=20, lr=1e-2
        )

        final_acc = history["val_acc"][-1]

        # Model should improve (even if not perfect on random val data)
        assert history["train_loss"][-1] < history["train_loss"][0]
        assert len(history["train_loss"]) == 20

    def test_history_has_correct_keys(self):
        key = jax.random.PRNGKey(0)
        params = init_params(key, 5, 10, 2)
        train_data = make_synthetic_dataset(key, 16, 5, 2, 8)
        val_data = (jax.random.normal(key, (8, 5)),
                    jax.random.randint(key, (8,), 0, 2))

        _, _, history = train_model(
            params, train_data, val_data, 2, epochs=3, lr=1e-3
        )

        assert "train_loss" in history
        assert "val_loss" in history
        assert "val_acc" in history
        assert len(history["train_loss"]) == 3

Testing on Different Devices and Precision

JAX code can run on CPU, GPU, and TPU. Your tests should be aware of device differences and precision implications. It's important to parameterize tests across devices and precision settings where relevant.

# tests/test_devices.py
import jax
import jax.numpy as jnp
import numpy as np
import pytest


@pytest.fixture(params=[
    ("float32", 1e-5),
    ("float64", 1e-10),
])
def precision(request):
    dtype, tol = request.param
    if dtype == "float64" and not jax.config.read("jax_enable_x64"):
        pytest.skip("64-bit precision not enabled")
    return getattr(jnp, dtype), tol


def test_precision_consistency(precision):
    dtype, tol = precision
    x = jnp.arange(10, dtype=dtype)
    result = jnp.sum(x ** 2)
    expected = jnp.array(285.0, dtype=dtype)
    np.testing.assert_allclose(result, expected, rtol=tol)


def test_device_placement():
    """Verify arrays