JAX ErrorsÂ¶
This page lists a few of the errors you might encounter when using JAX, along with representative examples of how one might fix them.

class
jax.errors.
ConcretizationTypeError
(tracer, context='')[source]Â¶ This error occurs when a JAX Tracer object is used in a context where a concrete value is required. In some situations, it can be easily fixed by marking problematic values as static; in others, it may indicate that your program is doing operations that are not directly supported by JAXâ€™s JIT compilation model.
 Traced value where static value is expected
One common cause of this error is using a traced value where a static value is required. For example:
>>> from jax import jit, partial >>> import jax.numpy as jnp >>> @jit ... def func(x, axis): ... return x.min(axis)
>>> func(jnp.arange(4), 0) Traceback (most recent call last): ... ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: axis argument to jnp.min().
This can often be fixed by marking the problematic value as static:
>>> @partial(jit, static_argnums=1) ... def func(x, axis): ... return x.min(axis) >>> func(jnp.arange(4), 0) DeviceArray(0, dtype=int32)
 Traced value used in control flow
Another case where this often arises is when a traced value is used in Python control flow. For example:
>>> @jit ... def func(x, y): ... return x if x.sum() < y.sum() else y >>> func(jnp.ones(4), jnp.zeros(4)) Traceback (most recent call last): ... ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: The problem arose with the `bool` function.
In this case, marking the problematic traced quantity as static is not an option, because it is derived from traced inputs. But you can make progress by reexpressing this if statement in terms of
jax.numpy.where()
:>>> @jit ... def func(x, y): ... return jnp.where(x.sum() < y.sum(), x, y) >>> func(jnp.ones(4), jnp.zeros(4)) DeviceArray([0., 0., 0., 0.], dtype=float32)
For more complicated control flow including loops, see Control flow operators.
 Shape depends on Traced Value
Such an error may also arise when a shape in your JITcompiled computation depends on the values within a traced quantity. For example:
>>> @jit ... def func(x): ... return jnp.where(x < 0) >>> func(jnp.arange(4)) Traceback (most recent call last): ... ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: The error arose in jnp.nonzero.
This is an example of an operation that is incompatible with JAXâ€™s JIT compilation model, which requires array sizes to be known at compiletime. Here the size of the returned array depends on the contents of x, and such code cannot be JIT compiled.
In many cases it is possible to work around this by modifying the logic used in the function; for example here is code with a similar issue:
>>> @jit ... def func(x): ... indices = jnp.where(x > 1) ... return x[indices].sum() >>> func(jnp.arange(4)) Traceback (most recent call last): ... ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: The error arose in jnp.nonzero.
And here is how you might express the same operation in a way that avoids creation of a dynamicallysized index array:
>>> @jit ... def func(x): ... return jnp.where(x > 1, x, 0).sum() >>> func(jnp.arange(4)) DeviceArray(5, dtype=int32)
To understand more subtleties having to do with tracers vs. regular values, and concrete vs. abstract values, you may want to read Different kinds of JAX values.
 Parameters
tracer (
Tracer
) â€“context (
str
) â€“

class
jax.errors.
NonConcreteBooleanIndexError
(tracer)[source]Â¶ This error occurs when a program attempts to use nonconcrete boolean indices in a traced indexing operation. Under JIT compilation, JAX arrays must have static shapes (i.e. shapes that are known at compiletime) and so boolean masks must be used carefully. Some logic implemented via boolean masking is simply not possible under JAXâ€™s JIT compilation model; in other cases, the logic can be reexpressed in a JITcompatible way, often using the threeargument version of
where()
.Following are a few examples of when this error might arise.
 Constructing arrays via boolean masking
This most commonly arises when attempting to create an array via a boolean mask within a JIT context. For example:
>>> import jax >>> import jax.numpy as jnp >>> @jax.jit ... def positive_values(x): ... return x[x > 0] >>> positive_values(jnp.arange(5, 5)) Traceback (most recent call last): ... NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10])
This function is attempting to return only the positive values in the input array; the size of this returned array cannot be determined at compiletime unless x is marked as static, and so operations like this cannot be performed under JIT compilation.
 Reexpressible Boolean Logic
Although creating dynamically sized arrays is not supported directly, in many cases it is possible to reexpress the logic of the computation in terms of a JITcompatible operation. For example, here is another function that fails under JIT for the same reason:
>>> @jax.jit ... def sum_of_positive(x): ... return x[x > 0].sum() >>> sum_of_positive(jnp.arange(5, 5)) Traceback (most recent call last): ... NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[10])
In this case, however, the problematic array is only an intermediate value, and we can instead express the same logic in terms of the JITcompatible threeargument version of
jax.numpy.where()
:>>> @jax.jit ... def sum_of_positive(x): ... return jnp.where(x > 0, x, 0).sum() >>> sum_of_positive(jnp.arange(5, 5)) DeviceArray(10, dtype=int32)
This pattern of replacing boolean masking with threeargument
where()
is a common solution to this sort of problem. Boolean indices in
jax.ops
The other situation where this error often arises is when using boolean indices within functions in
jax.ops
, such asjax.ops.index_update()
. Here is a simple example:>>> @jax.jit ... def manual_clip(x): ... return jax.ops.index_update(x, x < 0, 0) >>> manual_clip(jnp.arange(2, 2)) Traceback (most recent call last): ... NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[4])
This function is attempting to set values smaller than zero to a scalar fill value. As above, this can be addressed by reexpressing the logic in terms of
where()
:>>> @jax.jit ... def manual_clip(x): ... return jnp.where(x < 0, 0, x) >>> manual_clip(jnp.arange(2, 2)) DeviceArray([0, 0, 0, 1], dtype=int32)
These operations also commonly are written in terms of the Syntactic sugar for indexed update operators; for example, this is syntactic sugar for
index_mul()
, and fails under JIT:>>> @jax.jit ... def manual_abs(x): ... return x.at[x < 0].mul(1) >>> manual_abs(jnp.arange(2, 2)) Traceback (most recent call last): ... NonConcreteBooleanIndexError: Array boolean indices must be concrete: ShapedArray(bool[4])
As above, the solution is to reexpress this in terms of
where()
:>>> @jax.jit ... def manual_abs(x): ... return jnp.where(x < 0, x * 1, x) >>> manual_abs(jnp.arange(2, 2)) DeviceArray([2, 1, 0, 1], dtype=int32)
 Parameters
tracer (
Tracer
) â€“

class
jax.errors.
TracerArrayConversionError
(tracer)[source]Â¶ This error occurs when a program attempts to convert a JAX Tracer object into a standard NumPy array. It typically occurs in one of a few situations.
 Using numpy rather than jax.numpy functions
This error can occur when a JAX Tracer object is passed to a raw numpy function, or a method on a numpy.ndarray object. For example:
>>> from jax import jit, partial >>> import numpy as np >>> import jax.numpy as jnp >>> @jit ... def func(x): ... return np.sin(x) >>> func(jnp.arange(4)) Traceback (most recent call last): ... TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object
In this case, check that you are using jax.numpy methods rather than numpy methods:
>>> @jit ... def func(x): ... return jnp.sin(x) >>> func(jnp.arange(4)) DeviceArray([0. , 0.84147096, 0.9092974 , 0.14112 ], dtype=float32)
 Indexing a numpy array with a tracer
If this error arises on a line that involves array indexing, it may be that the array being indexed x is a raw numpy.ndarray while the indices idx are traced. For example:
>>> x = np.arange(10) >>> @jit ... def func(i): ... return x[i] >>> func(0) Traceback (most recent call last): ... TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object
Depending on the context, you may fix this by converting the numpy array into a JAX array:
>>> @jit ... def func(i): ... return jnp.asarray(x)[i] >>> func(0) DeviceArray(0, dtype=int32)
or by declaring the index as a static argument:
>>> @partial(jit, static_argnums=(0,)) ... def func(i): ... return x[i] >>> func(0) DeviceArray(0, dtype=int32)
To understand more subtleties having to do with tracers vs. regular values, and concrete vs. abstract values, you may want to read Different kinds of JAX values.
 Parameters
tracer (
Tracer
) â€“

class
jax.errors.
TracerIntegerConversionError
(tracer)[source]Â¶ This error can occur when a JAX Tracer object is used in a context where a Python integer is expected. It typically occurs in a few situations.
 Passing a tracer in place of an integer
This error can occur if you attempt to pass a tracer to a function that requires an integer argument; for example:
>>> from jax import jit, partial >>> import numpy as np >>> @jit ... def func(x, axis): ... return np.split(x, 2, axis) >>> func(np.arange(4), 0) Traceback (most recent call last): ... TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object
When this happens, the solution is often to mark the problematic argument as static:
>>> @partial(jit, static_argnums=1) ... def func(x, axis): ... return np.split(x, 2, axis) >>> func(np.arange(10), 0) [DeviceArray([0, 1, 2, 3, 4], dtype=int32), DeviceArray([5, 6, 7, 8, 9], dtype=int32)]
An alternative is to apply the transformation to a closure that encapsulates the arguments to be protected, either manually as below or by using
functools.partial()
:>>> jit(lambda arr: np.split(arr, 2, 0))(np.arange(4)) [DeviceArray([0, 1], dtype=int32), DeviceArray([2, 3], dtype=int32)]
Note a new closure is created at every invocation, which defeats the compilation caching mechanism, which is why static_argnums is preferred.
 Indexing a list with a Tracer
This error can occur if you attempt to index a Python list with a traced quantity. For example:
>>> import jax.numpy as jnp >>> from jax import jit, partial >>> L = [1, 2, 3] >>> @jit ... def func(i): ... return L[i] >>> func(0) Traceback (most recent call last): ... TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object
Depending on the context, you can generally fix this either by converting the list to a JAX array:
>>> @jit ... def func(i): ... return jnp.array(L)[i] >>> func(0) DeviceArray(1, dtype=int32)
or by declaring the index as a static argument:
>>> @partial(jit, static_argnums=0) ... def func(i): ... return L[i] >>> func(0) DeviceArray(1, dtype=int32)
To understand more subtleties having to do with tracers vs. regular values, and concrete vs. abstract values, you may want to read Different kinds of JAX values.
 Parameters
tracer (
Tracer
) â€“