jax-ml/jax · error · TypeError

Value of type {type(self)} is not indexable.

Error message

Value of type {type(self)} is not indexable.

What it means

Raised by JAXTracer.__setitem__ when item assignment (x[idx] = value) is attempted on a Tracer whose aval does not support mutation (arrays in JAX are immutable). JAX arrays and tracers cannot be modified in place, so __setitem__ is blocked with a TypeError; the comment notes ShapedArray provides a more detailed message for the common array case.

Source

Thrown at jax/_src/core.py:1162

    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.")
    return self.aval._array_module(self, types)

  def __getattr__(self, name):
    # if the aval property raises an AttributeError, gets caught here
    assert not config.enable_checks.value or name != "aval"

    # These must raise AttributeError in the base class for backward compatibility.
    # TODO(jakevdp): can we change this and make them raise NotImplementedError instead?
    if name in ["block_until_ready", "copy_to_host_async"]:
      raise AttributeError(
        f"The '{name}' method is not available on {self._error_repr()}."
        f"{self._origin_msg()}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use jax.numpy .at: x = x.at[i].set(v), x = x.at[i].add(v), x = x.at[mask].set(v) — functional updates that work under tracing.
  2. Restructure accumulation into carried state in jax.lax.scan/fori_loop (return new arrays each iteration).
  3. Use jnp.where for mask-based assignment instead of boolean-mask setitem.
  4. Allocate and fully construct new arrays (jnp.zeros(...).at[...].set(...)) rather than mutating inputs.

Example fix

# before
@jax.jit
def f(x):
    x[0] = 1.0        # TypeError: not indexable / arrays are immutable
    return x

# after
@jax.jit
def f(x):
    return x.at[0].set(1.0)
Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp
def safe_set(x, idx, v):
    return x.at[idx].set(v)   # works for both concrete arrays and tracers

Type guard

def is_mutable_host_array(x) -> bool:
    import numpy as np
    return isinstance(x, np.ndarray)

Prevention

When it happens

Trigger: In-place updates like x[0] = 1.0 or x[i, j] += 1 on a traced array inside jit/grad/vmap; initializing an output buffer with loop writes inside a jitted function; adapting NumPy mutation-style code without rewrite.

Common situations: Porting NumPy algorithms (Gauss-Seidel sweeps, buffer accumulation, mask assignment) directly into @jax.jit; trying to fill a preallocated array inside lax.scan bodies or vmap-mutated state; accumulation patterns like grads[mask] = 0.

Related errors

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