the museum
jax wing

serves the jax path · every error reproduced on this machine (jax 0.4.38, CPU, 2026-07-27), fixes run and proven · all wings

Every JAX surprise, met as an exhibit.

Six failures every JAX engineer meets in the first month, each one a chapter invariant wearing an error message. The snippet ran and failed; the error is verbatim (paths shortened); the fix ran and passed. Each caption names the rule, and the jax path teaches the rule before the error can find you.

exhibit
01/06

Writing into an array like it is memory

Chapter 01's first invariant, met as an error message. A jnp array is a value, not a place, so there is nothing to write into. .at[...].set expresses the same intent as a pure function, and XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → reuses the old buffer when the old value is dead.

the failing code

import jax.numpy as jnp

x = jnp.zeros(3)
x[0] = 1.0  # NumPy habit: update the array in place

the error, verbatimTypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html

the fix

import jax.numpy as jnp

x = jnp.zeros(3)
x = x.at[0].set(1.0)  # a new array; under jit the copy is usually elided

exhibit
02/06

Branching on a value the trace cannot see

Tracing runs your Python once with stand-ins that carry shape and dtype but no values, so a Python if on data has nothing to decide with. lax.cond stages both branches into the program and picks at run time, which is what you meant.

the failing code

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
    if x.sum() > 0:  # the tracer has no value to branch on
        return x
    return -x

f(jnp.ones(3))

the error, verbatimTracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]. The error occurred while tracing the function f at <string>:5 for jit. This concrete value was not available in Python because it depends on the value of the argument x. See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

the fix

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
    return jax.lax.cond(x.sum() > 0, lambda v: v, lambda v: -v, x)

f(jnp.ones(3))

exhibit
03/06

A boolean mask whose output shape depends on data

Shapes are fixed at trace time, and x[x > 0] has a length only the data knows. jnp.where keeps the shape and masks instead; when you truly need compaction, do it outside jit or with a static bound and padding.

the failing code

import jax
import jax.numpy as jnp

@jax.jit
def positives(x):
    return x[x > 0]  # output length depends on the values

positives(jnp.array([-1.0, 2.0, 3.0]))

the error, verbatimNonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[3]) See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError

the fix

import jax
import jax.numpy as jnp

@jax.jit
def positives(x):
    return jnp.where(x > 0, x, 0.0)  # same shape out, masked values

positives(jnp.array([-1.0, 2.0, 3.0]))

exhibit
04/06

Asking grad for the derivative of a vector

Reverse mode pulls one cotangent backward, so grad is defined for scalar outputs. Reduce to a scalar if you meant a loss; ask jacrev or jacfwd if you meant the whole Jacobian, and pick by the Jacobian's aspect ratio.

the failing code

import jax
import jax.numpy as jnp

g = jax.grad(lambda x: x * 2.0)(jnp.ones(3))

the error, verbatimTypeError: Gradient only defined for scalar-output functions. Output had shape: (3,).

the fix

import jax
import jax.numpy as jnp

g = jax.grad(lambda x: jnp.sum(x * 2.0))(jnp.ones(3))
J = jax.jacrev(lambda x: x * 2.0)(jnp.ones(3))  # the full Jacobian, if that was the ask

exhibit
05/06

vmap told to batch an argument that has no batch axis

in_axes defaults to 0 for every argument, and b has no axis of size 32 to give. Say which arguments carry the batch (0) and which ride along whole (None); the error names the sizes it saw, and they disagree for a reason.

the failing code

import jax
import jax.numpy as jnp

def dot(a, b):
    return a @ b

a = jnp.ones((32, 8))
b = jnp.ones(8)
jax.vmap(dot)(a, b)  # default in_axes=0 tries to split b too

the error, verbatimValueError: vmap got inconsistent sizes for array axes to be mapped: * one axis had size 32: axis 0 of argument a of type float32[32,8]; * one axis had size 8: axis 0 of argument b of type float32[8]

the fix

import jax
import jax.numpy as jnp

def dot(a, b):
    return a @ b

a = jnp.ones((32, 8))
b = jnp.ones(8)
out = jax.vmap(dot, in_axes=(0, None))(a, b)  # b broadcast to every row

exhibit
06/06

A static argument that cannot be a cache key

A static argument becomes part of the compilation cache key, so it must hash. Lists do not; tuples do. The chapter 03 rule applies twice here: statics are baked into the key, and every distinct value compiles its own executable.

the failing code

import jax
from functools import partial

@partial(jax.jit, static_argnames="dims")
def reduce_dims(x, dims):
    return x.sum(dims)

import jax.numpy as jnp
reduce_dims(jnp.ones((2, 3)), [0])  # a list cannot go in the cache key

the error, verbatimValueError: Non-hashable static arguments are not supported. An error occurred while trying to hash an object of type <class 'list'>, [0]. The error was: TypeError: unhashable type: 'list'

the fix

import jax
from functools import partial

@partial(jax.jit, static_argnames="dims")
def reduce_dims(x, dims):
    return x.sum(dims)

import jax.numpy as jnp
reduce_dims(jnp.ones((2, 3)), (0,))  # tuples hash; the key is the value itself