Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions autofit/jax/pytrees.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand All @@ -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))
Expand Down
Loading