Skip to content

fix: traced instance attributes never land in pytree aux data - #1441

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/jax-pytree-traced-aux-fix
Aug 1, 2026
Merged

fix: traced instance attributes never land in pytree aux data#1441
Jammy2211 merged 1 commit into
mainfrom
feature/jax-pytree-traced-aux-fix

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Part of PyAutoLabs/PyAutoLens#678 phase B — one of the two defects that killed the cluster image_plane_solved gradient cell on the A100s (the other is the PointSolver padded-row fix, PyAutoLabs/PyAutoLens#685).

Problem

_build_instance_pytree_funcs classifies attributes never declared on a walked model as constants → pytree aux data. Attributes derived inside __init__ from prior parameters (e.g. NFWMCRLudlowSph computing kappa_s/scale_radius from a free mass_at_200) are exactly such attributes, and under a trace their values are tracers. Aux crosses transformation boundaries as raw Python references, so a custom_jvp rule's inner jax.jvp receives the instance with a stale LinearizeTracer in scale_radius and raises UnexpectedTracerError the 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), isinstance dispatch) is untouched.

Verification

🤖 Generated with Claude Code

https://claude.ai/code/session_01DnTmLoJjJgMTze5uAbg1Jd

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
@Jammy2211
Jammy2211 merged commit 67c4090 into main Aug 1, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/jax-pytree-traced-aux-fix branch August 1, 2026 12:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant