jax-ml/jax · error · TypeError

Value of type {type(self)} is not convertible to integer ind

Error message

Value of type {type(self)} is not convertible to integer index.

What it means

Raised by JAXTracer.__index__ when a JAX Tracer is used where Python requires a concrete integer index (operator.index), e.g. as a list index, range argument, or tuple size. Traced values have no concrete integer value on the host, so they cannot serve as Python-level indices. JAX raises this TypeError because such use would break the trace being recorded.

Source

Thrown at jax/_src/core.py:1150

  def __hex__(self):
    if is_concrete(self): return hex(self.to_concrete_value())
    check_integer_conversion(self)
    if not hasattr(self.aval, "_hex"):
      raise TypeError(f"Value of type {type(self)} is not convertible to hex.")
    return self.aval._hex(self)

  def __oct__(self):
    if is_concrete(self): return oct(self.to_concrete_value())
    check_integer_conversion(self)
    if not hasattr(self.aval, "_oct"):
      raise TypeError(f"Value of type {type(self)} is not convertible to oct.")
    return self.aval._oct(self)

  def __index__(self):
    if is_concrete(self): return operator.index(self.to_concrete_value())
    check_integer_conversion(self)
    if not hasattr(self.aval, "_index"):
      raise TypeError(f"Value of type {type(self)} is not convertible to integer index.")
    return self.aval._index(self)

  # raises a useful error on attempts to pickle a Tracer.
  def __reduce__(self):
    raise ConcretizationTypeError(
      self, ("The error occurred in the __reduce__ method, which may "
             "indicate an attempt to serialize/pickle a traced value."))

  # raises the better error message from ShapedArray
  def __setitem__(self, key, value):
    if not hasattr(self.aval, "_setitem"):
      raise TypeError(f"Value of type {type(self)} is not indexable.")
    return self.aval._setitem(self, key, value)

  # NumPy also only looks up special methods on classes.
  def __array_module__(self, types):
    if not hasattr(self.aval, "_array_module"):
      raise TypeError(f"Value of type {type(self)} is not compatible with the Array API.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert the Python list/tuple to a jax.numpy array first: jnp.asarray(lst)[tracer] uses JAX dynamic indexing and works under jit.
  2. If the index is actually constant, mark it static with functools.partial(jit, static_argnums=...) so it arrives as a plain int.
  3. For data-dependent loops, use jax.lax.fori_loop / scan / cond instead of Python for/if on traced values.
  4. If you truly need the concrete value, restructure so the index is computed outside the traced function.

Example fix

# before
@jax.jit
def f(table, i):
    return table[i]        # table is a Python list, i a Tracer

# after
table_arr = jnp.asarray(table)
@jax.jit
def f(i):
    return table_arr[i]    # JAX array supports dynamic (traced) indices
Defensive patterns

Strategy: type-guard

Validate before calling

import jax
import jax.numpy as jnp
def lookup(table, i):
    if isinstance(i, jax.core.Tracer) or isinstance(table, (list, tuple)):
        return jnp.asarray(table)[i]   # dynamic index path
    return table[i]

Type guard

import jax
def needs_dynamic_index(seq, i) -> bool:
    return isinstance(seq, (list, tuple)) and isinstance(i, jax.core.Tracer)

Prevention

When it happens

Trigger: Indexing a Python list/tuple with a traced integer (lst[tracer]) inside jit/grad; range(tracer) or np.arange(tracer_count) style code; slicing host-side sequences with traced bounds; using a traced value as a shape dimension in host-side code.

Common situations: Data-dependent control flow like data = dataset[batch_size_tracer] inside jitted functions; computing loop counts under vmap/scan; mixing Python collections with traced scalars instead of jax.numpy arrays.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/b18bbcda3f9c30da. Report an issue: GitHub.