jax-ml/jax · error · ValueError

Triggering __jax_array__() during abstractification is no lo

Error message

Triggering __jax_array__() during abstractification is no longer supported. To avoid this error, either explicitly convert your object using jax.numpy.array(), or register your object as a pytree. 

What it means

JAX no longer implicitly calls __jax_array__() while abstractifying an object (converting it to a JAX abstract value during tracing). If your custom class defines __jax_array__ but is not registered as a pytree, passing its instances into JAX functions raises this ValueError.

Source

Thrown at jax/_src/core.py:1983

# We have two flavors of abstractification APIs here which each used to have
# their own separate implementation. Now they're effectively the same, with the
# following differences:
#
# - typeof returns avals for valid array-like objects, including tracers.
# - shaped_abstractify is like typeof, but also accepts duck-typed arrays.
#

def shaped_abstractify(x):
  typ = type(x)
  if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path
    return aval_fn(x)
  for t in typ.__mro__[1:]:
    if (aval_fn := pytype_aval_mappings.get(t)):
      return aval_fn(x)
  if isinstance(x, AbstractValue):
    return x
  if getattr(x, '__jax_array__', None) is not None:
    raise ValueError(
        'Triggering __jax_array__() during abstractification is no longer'
        ' supported. To avoid this error, either explicitly convert your object'
        ' using jax.numpy.array(), or register your object as a pytree.'
    )
  if hasattr(x, 'dtype'):
    aval = ShapedArray(
        np.shape(x),
        dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True),
        weak_type=getattr(x, "weak_type", False),
    )
    return update_aval_with_sharding(aval, getattr(x, 'sharding', None))
  raise TypeError(
      f"Cannot interpret value of type {typ} as an abstract array; it "
      "does not have a dtype attribute")


# TODO(phawkins): the return type should be AbstractValue.
def typeof(x: Any) -> Any:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert explicitly with jax.numpy.array(obj) before passing it into JAX
2. Register the class as a pytree with jax.tree_util.register_pytree_node (or @jax.tree_util.register_pytree_dataclass / struct.dataclass)
3. Add a to-jax conversion in the wrapper's __init__ so the stored value is already a jnp.ndarray

Example fix

// before
class Wrap:
    def __init__(self, a): self.a = a
    def __jax_array__(self): return self.a
jax.jit(fn)(Wrap(x))  # ValueError

// after
jax.jit(fn)(jnp.array(Wrap(x)))
# or register as pytree
jax.tree_util.register_pytree_node(Wrap, lambda w: ((w.a,), None), lambda n, c: Wrap(c[0]))

import jax
def jax_ready(obj):
    return hasattr(obj, 'dtype') or jax.tree_util.all_leaves([obj]) or hasattr(obj, 'tree_flatten')
# or simply pre-convert: obj = jnp.array(obj)

lambda x: not hasattr(type(x), '__jax_array__') or jax.tree_util.all_leaves([x])

try:
    jax.jit(fn)(obj)
except ValueError as e:
    if '__jax_array__' in str(e):
        jax.jit(fn)(jnp.array(obj))
    else:
        raise

When it happens

Trigger: Passing an instance of a user-defined class with a __jax_array__ method directly as an argument to jit/vmap/grad-wrapped functions or jnp functions, without registering it via jax.tree_util.register_pytree_node or converting with jax.numpy.array().

Common situations: Upgrading JAX versions where implicit __jax_array__ conversion was removed; wrapping arrays in wrapper dataclasses for metadata; integrating third-party array-like objects (e.g., custom tensor wrappers).

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• The return value of the policies should be a boolean. Got: {

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