From 2da81578e5c360deeabc3ea91ad834552bc382f9 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 13 Jul 2026 19:26:22 +0100 Subject: [PATCH] fix: don't register JAX-leaf types (int) as instance pytree nodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit register_model walked every Model.cls into a pytree registration, including the af.Model(int) sub-models autofit auto-creates for int-typed constructor args (e.g. a shapelet's n/m). Registering builtins.int made JAX call the instance flatten _partition(5) on every int leaf, and vars(5) raised 'TypeError: vars() argument must have __dict__ attribute', crashing jax.vmap(jax.jit(...)) for any JAX shapelet fit (release-validation run 29266305445). Gate registration on _instances_carry_dict(cls): only classes whose instances carry a __dict__ (real model components) are registered; native JAX leaves (int/float/str/...) and purely __slots__ classes are left as opaque leaves. Producer-side fix — no silent fallback inside _partition. Fixes #1365 Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01TDCsgNCXacXqiWAPTswWCt --- autofit/jax/pytrees.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/autofit/jax/pytrees.py b/autofit/jax/pytrees.py index cd9cdfde1..5cd2bc8bc 100644 --- a/autofit/jax/pytrees.py +++ b/autofit/jax/pytrees.py @@ -15,6 +15,27 @@ _CLASS_CONSTRUCTOR_ARGS: dict = {} +def _instances_carry_dict(cls) -> bool: + """Whether instances of ``cls`` have a ``__dict__``. + + Only such classes may be registered as instance pytree nodes: the flatten + built by ``_build_instance_pytree_funcs`` reads ``vars(instance)``, which + raises ``TypeError`` on an instance without a ``__dict__``. + + This excludes native JAX leaf types (``int``, ``float``, ``str``, ``bool``, + ``NoneType``, ``tuple``, ...) — which JAX already handles correctly as + opaque leaves — and any purely ``__slots__`` class. ``register_model`` + reaches these via ``af.Model(int)`` sub-models (e.g. a ``ShapeletPolar``'s + integer ``n``/``m`` constructor args): without this guard ``int`` gets + registered, JAX then calls the instance-flatten on every ``int`` leaf, and + ``vars(5)`` crashes the whole ``jax.jit``/``jax.vmap`` flatten. + + A class defines a ``__dict__`` for its instances iff some class in its MRO + carries the ``__dict__`` descriptor. + """ + return any("__dict__" in vars(klass) for klass in cls.__mro__) + + def enable_pytrees() -> bool: """Register PyAutoFit model and prior classes as JAX pytree nodes. @@ -92,7 +113,14 @@ def _walk(node): _CLASS_CONSTRUCTOR_ARGS.setdefault( cls, tuple(node.constructor_argument_names) ) - if cls not in _REGISTERED_INSTANCE_CLASSES: + # Only register classes whose instances carry a ``__dict__`` — the + # instance flatten reads ``vars(instance)``. ``af.Model(int)`` + # sub-models (e.g. a shapelet's integer ``n``/``m`` args) would + # otherwise register ``int`` and crash the flatten on every int leaf. + if ( + cls not in _REGISTERED_INSTANCE_CLASSES + and _instances_carry_dict(cls) + ): flatten, unflatten = _build_instance_pytree_funcs(cls) try: register_pytree_node(cls, flatten, unflatten)