fix: traced instance attributes never land in pytree aux data - #1441
Merged
Conversation
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DnTmLoJjJgMTze5uAbg1Jd
This was referenced Aug 3, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of PyAutoLabs/PyAutoLens#678 phase B — one of the two defects that killed the cluster
image_plane_solvedgradient cell on the A100s (the other is the PointSolver padded-row fix, PyAutoLabs/PyAutoLens#685).Problem
_build_instance_pytree_funcsclassifies attributes never declared on a walked model as constants → pytree aux data. Attributes derived inside__init__from prior parameters (e.g.NFWMCRLudlowSphcomputingkappa_s/scale_radiusfrom a freemass_at_200) are exactly such attributes, and under a trace their values are tracers. Aux crosses transformation boundaries as raw Python references, so acustom_jvprule's innerjax.jvpreceives the instance with a staleLinearizeTracerinscale_radiusand raisesUnexpectedTracerErrorthe moment it is consumed (1.0 / self.scale_radius).Only gradient paths hit it (Nautilus never runs a JVP rule), and only models with derived-in-
__init__attributes (galaxy-scale Isothermal cells were unaffected) — which is why this survived until the first cluster MCR-halo gradient search.Fix
At flatten time, any attribute whose value is a JAX array or tracer is promoted to a dynamic child, whatever the classifier says — a traced value is never safe aux. Concrete Python values keep the constant behaviour, so control flow on constants (
sorted(..., key=lambda g: g.redshift),isinstancedispatch) is untouched.Verification
test_autofit/: 1641 passed, 2 skipped.autofit_workspace_test/scripts/jax_assertions/derived_attribute_leaves.py(JAX tests live there per test: delete jax-using unit tests (moved to autofit_workspace_test) #1247): fails withUnexpectedTracerErrorbefore this fix, passes after.value_and_gradcompletes after this fix (paired with fix: sanitize padded rows in the PointSolver implicit-diff rule for reverse mode PyAutoLens#685 for finite gradients).🤖 Generated with Claude Code
https://claude.ai/code/session_01DnTmLoJjJgMTze5uAbg1Jd