{"record":{"id":"f8fab4a42e449f23","repo":"jax-ml/jax","slug":"triggering-jax-array-during-abstractificatio","errorCode":null,"errorMessage":"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. ","messagePattern":"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\\. ","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":1983,"sourceCode":"# We have two flavors of abstractification APIs here which each used to have\n# their own separate implementation. Now they're effectively the same, with the\n# following differences:\n#\n# - typeof returns avals for valid array-like objects, including tracers.\n# - shaped_abstractify is like typeof, but also accepts duck-typed arrays.\n#\n\ndef shaped_abstractify(x):\n  typ = type(x)\n  if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path\n    return aval_fn(x)\n  for t in typ.__mro__[1:]:\n    if (aval_fn := pytype_aval_mappings.get(t)):\n      return aval_fn(x)\n  if isinstance(x, AbstractValue):\n    return x\n  if getattr(x, '__jax_array__', None) is not None:\n    raise ValueError(\n        'Triggering __jax_array__() during abstractification is no longer'\n        ' supported. To avoid this error, either explicitly convert your object'\n        ' using jax.numpy.array(), or register your object as a pytree.'\n    )\n  if hasattr(x, 'dtype'):\n    aval = ShapedArray(\n        np.shape(x),\n        dtypes.canonicalize_dtype(x.dtype, allow_extended_dtype=True),\n        weak_type=getattr(x, \"weak_type\", False),\n    )\n    return update_aval_with_sharding(aval, getattr(x, 'sharding', None))\n  raise TypeError(\n      f\"Cannot interpret value of type {typ} as an abstract array; it \"\n      \"does not have a dtype attribute\")\n\n\n# TODO(phawkins): the return type should be AbstractValue.\ndef typeof(x: Any) -> Any:","sourceCodeStart":1965,"sourceCodeEnd":2001,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L1965-L2001","documentation":"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.","triggerScenarios":"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().","commonSituations":"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).","solutions":["Convert explicitly with jax.numpy.array(obj) before passing it into JAX","Register the class as a pytree with jax.tree_util.register_pytree_node (or @jax.tree_util.register_pytree_dataclass / struct.dataclass)","Add a to-jax conversion in the wrapper's __init__ so the stored value is already a jnp.ndarray"],"exampleFix":"// before\nclass Wrap:\n    def __init__(self, a): self.a = a\n    def __jax_array__(self): return self.a\njax.jit(fn)(Wrap(x))  # ValueError\n\n// after\njax.jit(fn)(jnp.array(Wrap(x)))\n# or register as pytree\njax.tree_util.register_pytree_node(Wrap, lambda w: ((w.a,), None), lambda n, c: Wrap(c[0]))","handlingStrategy":"validation","validationCode":"import jax\ndef jax_ready(obj):\n    return hasattr(obj, 'dtype') or jax.tree_util.all_leaves([obj]) or hasattr(obj, 'tree_flatten')\n# or simply pre-convert: obj = jnp.array(obj)","typeGuard":"lambda x: not hasattr(type(x), '__jax_array__') or jax.tree_util.all_leaves([x])","tryCatchPattern":"try:\n    jax.jit(fn)(obj)\nexcept ValueError as e:\n    if '__jax_array__' in str(e):\n        jax.jit(fn)(jnp.array(obj))\n    else:\n        raise","preventionTips":["Register custom array wrappers as pytrees at import time","Convert at the boundary with jnp.array instead of relying on __jax_array__","Pin/audit JAX version changes that alter abstractification behavior"],"tags":["jax","pytree","jax-array","custom-class","tracing"],"backgroundTag":"jax-pytree-registration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}