From 9283d2d22c3a425ac12f3063e6f6931b2b16a439 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 1 Aug 2026 13:21:03 +0100 Subject: [PATCH] fix: traced instance attributes never land in pytree aux data Attributes derived inside __init__ from prior parameters (an NFWMCRLudlowSph computing scale_radius from a free mass_at_200) are unknown to the instance pytree classifier and defaulted to constant aux. Under a trace those values ARE tracers; aux survives flatten as raw Python references and re-enters nested traces (a custom_jvp rule's inner jax.jvp) as stale tracers, raising UnexpectedTracerError (PyAutoLens#678 phase B cluster gradient cells). Flatten now promotes any attribute whose value is a JAX array or tracer to a dynamic child. Concrete Python values keep the constant/aux behaviour, so control flow reading constants (sorted by redshift, isinstance dispatch) is untouched. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01DnTmLoJjJgMTze5uAbg1Jd --- autofit/jax/pytrees.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/autofit/jax/pytrees.py b/autofit/jax/pytrees.py index 105f1f2c9..7a705396e 100644 --- a/autofit/jax/pytrees.py +++ b/autofit/jax/pytrees.py @@ -166,11 +166,27 @@ def _build_instance_pytree_funcs(cls): Classification is read from the shared ``_CLASS_FIELD_CLASSIFIERS`` dict, which is updated by every ``register_model`` call. Attributes unknown to the classifier (never declared on any walked model) default to constant — - safer than tracing an unknown object. + safer than tracing an unknown object — with one override: an attribute + whose *value* is a JAX array or tracer is always a dynamic child, + whatever the classifier says. Such values arise from attributes derived + inside ``__init__`` from traced parameters (e.g. an ``NFWMCRLudlowSph`` + computing ``scale_radius`` from a free ``mass_at_200``); as aux data they + survive the flatten as raw Python references and re-enter nested traces + (a ``custom_jvp`` rule's inner jvp) as stale tracers, raising + ``UnexpectedTracerError``. A traced value is never safe aux. """ constructor_args = _CLASS_CONSTRUCTOR_ARGS.get(cls, ()) constructor_arg_set = set(constructor_args) + def _is_jax_value(value): + import jax + + if isinstance(value, (jax.Array, jax.core.Tracer)): + return True + if isinstance(value, (tuple, list)): + return any(_is_jax_value(v) for v in value) + return False + def _partition(instance): classifier = _CLASS_FIELD_CLASSIFIERS.get(cls, {}) ctor_dyn: list = [] @@ -180,7 +196,7 @@ def _partition(instance): for name, value in vars(instance).items(): if name.startswith("_") or name in ("cls", "id"): continue - is_dynamic = classifier.get(name, False) + is_dynamic = classifier.get(name, False) or _is_jax_value(value) in_ctor = name in constructor_arg_set if in_ctor and is_dynamic: ctor_dyn.append((name, value))