What a tracer refuses

Before the refusals, the permissions, because they are wider than people expect. Shape, length, dtype and rank all answer from the aval. Arithmetic on those answers is ordinary Python arithmetic, so jnp.zeros(len(x) * 2) builds a (6,) array inside a trace with nothing traced about it at all.

Anything that produces a new traced value is fine too. x.astype(jnp.int32) records a convert equation, x.at[0].set(9) records a scatter, and both hand back another tracer for the next line to use. The recording grows and nothing needs a number.

The line between the two groups is one question: does this call need to know a value right now, in Python, before the trace can continue? Shape arithmetic does not. A branch does.

import jax
import jax.numpy as jnp

@jax.jit
def free(x):
    print(x.shape, len(x), x.dtype, x.ndim)
    print(jnp.zeros(len(x) * 2).shape)
    return x.astype(jnp.int32).at[0].set(9)

print(free(jnp.ones(3)))

# (3,) 3 float32 1
# (6,)
# [9 1 1]

The six lines below all ask a tracer for a value, and JAX distinguishes them by which Python method got called. An if reaches __bool__ and raises TracerBoolConversionError. A float() or an .item() reaches the concretization path and raises ConcretizationTypeError. Handing a tracer to NumPy reaches __array__ and raises TracerArrayConversionError. Indexing a Python list with a tracer reaches __index__ and raises TracerIntegerConversionError.

Two entries in that run are worth pausing on. min(x[0], x[1]) never mentions a branch in your source, and it raises the branch error, because min compares its arguments and a comparison of tracers has to be resolved to a Python bool. Anything built on comparison lands in the same place: max, sorted with a key that compares traced values, a while condition, an assert.

The other one is .item(), which people reach for precisely when they want to look at a number. It shares a class with float() rather than getting one of its own, and the message says which method asked, so the class narrows the cause and the first line finishes the job.

The class of the error names the method that asked. The message names the line.

import jax
import jax.numpy as jnp
import numpy as np

def refuse(name, fn):
    try:
        jax.jit(fn)(jnp.ones(3))
    except Exception as e:
        print(name, type(e).__name__, '|', str(e).splitlines()[0])

refuse('if      ', lambda x: x if x.sum() > 0 else -x)
refuse('float() ', lambda x: float(x.sum()))
refuse('.item() ', lambda x: x.sum().item())
refuse('np.array', lambda x: np.asarray(x))
refuse('index   ', lambda x: [1.0, 2.0][x.sum().astype(int)])
refuse('min()   ', lambda x: min(x[0], x[1]))

# if       TracerBoolConversionError | Attempted boolean conversion of traced array with shape bool[].
# float()  ConcretizationTypeError | Abstract tracer value encountered where concrete value is expected: traced array with shape float32[]
# .item()  ConcretizationTypeError | Abstract tracer value encountered where concrete value is expected: traced array with shape float32[]
# np.array TracerArrayConversionError | The numpy.ndarray conversion method __array__() was called on traced array with shape float32[3]
# index    TracerIntegerConversionError | The __index__() method was called on traced array with shape int32[]
# min()    TracerBoolConversionError | Attempted boolean conversion of traced array with shape bool[].

Read one of those messages in full and it has four parts doing four jobs. The first line states the abstract-value problem and prints the aval. The second names the Python function that asked, float here, and offers the conversion that would have worked. The third points at the traced function and the line it was defined on. The fourth is a link to the error's own documentation page.

The third line is the one people skim and the one that saves the most time on a large program, because it names the function being traced rather than the frame you happened to be looking at. When a refusal fires four calls deep inside a library, that line is what tells you whose trace you are inside.

The message also states, plainly, why the value was unavailable: it depends on the value of the argument. That is the distinction chapter 3 turns into a decision about static_argnums, and reading it here is what makes that decision obvious later.

Abstract tracer value encountered where concrete value is expected: traced array with shape float32[]
The problem arose with the `float` function. If trying to convert the data type of a value, try using `x.astype(float)` or `jnp.array(x, float)` instead.
The error occurred while tracing the function scale at <string>:4 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.ConcretizationTypeError

# produced by:
# import jax
# import jax.numpy as jnp
#
# @jax.jit
# def scale(x):
#     return x / float(x.sum())
#
# try:
#     scale(jnp.ones(3))
# except jax.errors.ConcretizationTypeError as e:
#     print(e)

Not every failure inside a trace comes from the tracer error family, and the ones that do not are the ones that waste an afternoon. bool(x) on a three-element array raises a plain ValueError about ambiguous truth values, the same sentence NumPy raises, because the shape check runs first and the shape is known. Rank is not the problem the trace has, so that check answers before the value check ever runs.

Formatting is the second. Putting a traced scalar in an f-string raises TypeError: unsupported format string passed to DynamicJaxprTracer.__format__, which is the only refusal in this lesson that names the tracer class in the message and gives no advice at all.

The third is the one that shows up in real code most often. Ask for an array whose size comes from a traced value and the shape machinery refuses before any tracer method is reached: Shapes must be 1D sequences of concrete values of integer type, with the tracer printed inside the tuple it was found in. The museum's boolean-mask exhibit is the same wall approached from the other side.

import jax
import jax.numpy as jnp

def fails(fn):
    try:
        jax.jit(fn)(jnp.ones(3))
    except Exception as e:
        print(type(e).__name__, '|', str(e).splitlines()[0])

fails(lambda x: x if bool(x) else -x)
fails(lambda x: f"{x.sum():.2f}")
fails(lambda x: jnp.zeros(x.sum().astype(int)))

# ValueError | The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
# TypeError | unsupported format string passed to DynamicJaxprTracer.__format__
# TypeError | Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int32[])>with<DynamicJaxprTrace>,).

The four tracer error classes are not siblings in the way the messages suggest. TracerBoolConversionError is a subclass of ConcretizationTypeError, so catching the general one catches the branch failure too. TracerIntegerConversionError and TracerArrayConversionError are not: they sit directly under JAXTypeError, alongside ConcretizationTypeError rather than beneath it.

Everything in the tree bottoms out at Python's TypeError, UnexpectedTracerError from the next lesson included. A bare except TypeError around a traced call therefore swallows every refusal in this lesson, which is a good reason not to write one.

For a fixture that has to survive any refusal, jax.errors.JAXTypeError is the one class that covers all five without reaching outside JAX.

import jax

for name in ["TracerBoolConversionError", "TracerIntegerConversionError",
             "TracerArrayConversionError", "ConcretizationTypeError",
             "UnexpectedTracerError"]:
    cls = getattr(jax.errors, name)
    print(name, "<-", " <- ".join(b.__name__ for b in cls.__mro__[1:4]))

# TracerBoolConversionError <- ConcretizationTypeError <- JAXTypeError <- _JAXErrorMixin
# TracerIntegerConversionError <- JAXTypeError <- _JAXErrorMixin <- TypeError
# TracerArrayConversionError <- JAXTypeError <- _JAXErrorMixin <- TypeError
# ConcretizationTypeError <- JAXTypeError <- _JAXErrorMixin <- TypeError
# UnexpectedTracerError <- JAXTypeError <- _JAXErrorMixin <- TypeError