jax-ml/jax · error · TypeError

Value of type {type(self)} is not compatible with the Array

Error message

Value of type {type(self)} is not compatible with the Array API.

What it means

Raised by JaxTracer.__array_module__ when a library using the Python array API standard protocol asks a Tracer for its array module (via __array_module__(types)) and the tracer's aval does not implement _array_module. This protocol is how NumPy>=1.22 and array-api-compatible libraries decide whether they can dispatch NEP-35 style operations on an object; tracers cannot participate, so JAX refuses.

Source

Thrown at jax/_src/core.py:1168

      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()}")

    if name == 'sharding':
      raise AttributeError(
        f"The 'sharding' attribute is not available on {self._error_repr()}. "
        "To query sharding information on tracers, use `jax.typeof(x)`.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace the NumPy/SciPy call with its jax.numpy / jax.scipy equivalent inside the traced function.
2. If the operation cannot run under tracing, hoist it out: return the value from jit, do the NumPy work on the concrete host array, then re-enter JAX.
3. Mark genuinely static arguments via static_argnums/static_argnames so they arrive as normal arrays/objects.
4. For third-party numerics, check if the library supports JAX via the array API or a jax backend parameter.

Example fix

# before
@jax.jit
def f(x):
    return np.sqrt(x) + scipy.special.gammaln(x)  # probes __array_module__

# after
import jax.numpy as jnp
import jax.scipy as jsp
@jax.jit
def f(x):
    return jnp.sqrt(x) + jsp.special.gammaln(x)

import jax
def is_tracer(x) -> bool:
    return isinstance(x, jax.core.Tracer)
# route dispatch before calling numpy code
def dispatch(x):
    return (jax.numpy if is_tracer(x) else __import__('numpy'))

import jax
def is_tracer(x) -> bool:
    return isinstance(x, jax.core.Tracer)

When it happens

Trigger: Calling np.ndarray.__array_module__ dispatch paths on a tracer, e.g. numpy_function(jax_tracer) where NumPy inspects the object via the array module protocol; using array-api-compat or libraries like scipy that probe __array_module__; np.asarray-like coercion inside jit/grad where the tracer is passed to host code.

Common situations: Passing traced JAX values to NumPy/SciPy functions that don't understand NEP-35 dispatch; calling pandas/sklearn utilities on values captured inside a jitted function; version upgrades where NumPy started preferring __array_module__ over __array__ for protocol negotiation.

Related errors

• iteration over a 0-d array
• Value of type {type(self)} is not convertible to float.
• Value of type {type(self)} is not convertible to complex.
• Value of type {type(self)} is not convertible to hex.
• Value of type {type(self)} is not convertible to oct.

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