feat: re-key linear_light_profile_intensity_dict to survive jax.jit round-trip

[Jammy2211]

Overview

Make fit.linear_light_profile_intensity_dict survive a jax.jit round-trip so the visualization path works for any model containing a linear light profile (al.lp_linear.Gaussian, al.lmp_linear.GaussianGradient, etc.). Currently, Part 2 of modeling_visualization_jit.py only works with parametric Sersic — any LightProfileLinear dies with KeyError inside the visualizer callback because the dict is identity-keyed and pytree unflatten produces fresh objects.

Plan

• Attach a stable _pytree_token (monotonic counter) to every LightProfileLinear instance
• Override __hash__ / __eq__ on LightProfileLinear to use the token instead of id()
• Add __exclude_identifier_fields__ so the token never affects PyAutoFit's model unique_id
• Register LightProfileLinear as a pytree with no_flatten=("_pytree_token",) so the token rides as static aux data across JIT
• Revert Part 2 of modeling_visualization_jit.py to full MGE (parametric basis of GaussianGradient, NFWSph mass, MGE source) and verify subplot_fit.png lands with no KeyError

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary — hosts LightProfileLinear)
• PyAutoLens (pytree registration site)
• autolens_workspace_test (Part 2 revert)

Work Classification

Both (library + workspace)

Branch Survey

Repository	Current Branch	Dirty?
PyAutoGalaxy	main	clean
PyAutoLens	main	clean
autolens_workspace_test	main	clean

Suggested branch: feature/linear-light-profile-intensity-dict-pytree
Worktree root: ~/Code/PyAutoLabs-wt/linear-light-profile-intensity-dict-pytree/

Option chosen: B (stable identifier token)

Option B over A (structural path) and C (value eq/hash). A requires threading paths through linear_obj_func.light_profile_list construction — higher blast radius. C risks surprise collisions between profiles with identical params.

Implementation Steps

1. In PyAutoGalaxy/autogalaxy/profiles/light/linear/abstract.py, add to LightProfileLinear:

   import itertools
   _pytree_token_counter = itertools.count()
   __exclude_identifier_fields__ = ("_pytree_token",)
   
   def __init__(self, *args, **kwargs):
       super().__init__(*args, **kwargs)
       self._pytree_token = next(LightProfileLinear._pytree_token_counter)
   
   def __hash__(self):
       return self._pytree_token
   
   def __eq__(self, other):
       return isinstance(other, LightProfileLinear) and self._pytree_token == other._pytree_token

2. In PyAutoLens/autolens/imaging/model/analysis.py _register_fit_imaging_pytrees (after line 151): register_instance_pytree(LightProfileLinear, no_flatten=("_pytree_token",))
3. Revert Part 2 of autolens_workspace_test/scripts/imaging/modeling_visualization_jit.py to full MGE model (look at git log for the pre-split state).
4. Verify:
   • modeling_visualization_jit.py runs end-to-end, subplot_fit.png lands, no KeyError
   • multi_galaxy_mge.py regression identifier 5a3c480de681f6958048b22b3db8ecf9 unchanged (Option B audit confirmed identifier naturally filters _pytree_token because it's not an __init__ arg; __exclude_identifier_fields__ is belt-and-braces)
   • PyAutoGalaxy + PyAutoLens test suites green

Key Files

• PyAutoGalaxy/autogalaxy/profiles/light/linear/abstract.py — LightProfileLinear class + lp_instance_from at line 148
• PyAutoGalaxy/autogalaxy/abstract_fit.py:122 — linear_light_profile_intensity_dict construction
• PyAutoLens/autolens/imaging/model/analysis.py — _register_fit_imaging_pytrees
• autolens_workspace_test/scripts/imaging/modeling_visualization_jit.py — Part 2

Audits completed (de-risked Option B)

• Class structure: LightProfileLinear is a plain class, no dataclass/slots. MRO override of __hash__/__eq__ wins cleanly over GeometryProfile ancestors. Minor caveat: pickled instances retain pickled token.
• Call sites: the only dict that changes is the target one. LinearObj subclasses (Mapper, LightProfileLinearObjFuncList) are not LightProfileLinear; other identity-keyed dicts unaffected.
• Identifier: PyAutoFit/autofit/mapper/identifier.py:127-128 filters __dict__ to __init__ args only, so _pytree_token is naturally excluded. __exclude_identifier_fields__ is added explicitly for clarity + insurance against future code paths.

Original Prompt

Click to expand starting prompt

Make `fit.linear_light_profile_intensity_dict` survive a `jax.jit` round-trip so the visualization path works for any model containing a linear light profile. See `admin_jammy/prompt/issued/linear_light_profile_intensity_dict_pytree.md` for the full text.

[Jammy2211]

Library PR Created — 2026-04-19

PRs

• feat: stable pytree_token on LightProfileLinear for jax.jit round-trip PyAutoGalaxy#360
• fix: off-by-one in test__lp_linear_func_galaxy_dict_from exposed by LightProfileLinear eq #449

Status

Workspace changes needed (Part 2 of modeling_visualization_jit.py reverted to full MGE). Next: /ship_workspace

API Changes

LightProfileLinear gains pytree_token + token-based __hash__ / __eq__ so linear_light_profile_intensity_dict survives pytree unflatten. Previously structural __dict__ eq; now identity-equivalent (same object only).

[Jammy2211]

Shipped — 2026-04-19

PRs

• feat: stable pytree_token on LightProfileLinear for jax.jit round-trip PyAutoGalaxy#360 (LightProfileLinear pytree_token)
• fix: off-by-one in test__lp_linear_func_galaxy_dict_from exposed by LightProfileLinear eq #449 (off-by-one test fix exposed by new eq)
• test: revert modeling_visualization_jit Part 2 to full MGE autolens_workspace_test#34 (modeling_visualization_jit Part 2 → full MGE)

Summary

linear_light_profile_intensity_dict now survives jax.jit round-trip. Implemented Option B (stable identifier token) from the prompt: every LightProfileLinear instance gets a monotonic pytree_token at __init__; __hash__ and __eq__ route through the token; name has no leading underscore so PyAutoFit's register_model preserves it as attr_const across pytree flatten/unflatten; __exclude_identifier_fields__ keeps it out of the model unique_id hash. End-to-end validation: Nautilus runs full MGE model (GaussianGradient basis + NFWSph + MGE source), fit.png lands during quick-update, no KeyError.

Session Notes

• Option B over A/C: A (structural path key) requires threading paths through linear_obj_func.light_profile_list construction — higher blast radius. C (value eq/hash) risks surprise collisions when two linear profiles have identical params.
• Name matters: _pytree_token (leading underscore) would be silently skipped by register_model at PyAutoFit/autofit/jax/pytrees.py:153 — the token wouldn't ride pytree aux, breaking the fix. pytree_token (no underscore) is preserved via the default attr_const path.
• Identifier safety: Audited PyAutoFit/autofit/mapper/identifier.py:127-128 — filters __dict__ to __init__ args only, so pytree_token is naturally excluded. __exclude_identifier_fields__ added belt-and-braces. multi_galaxy_mge.py regression hash 5a3c480de681f6958048b22b3db8ecf9 confirmed unchanged.
• Pre-existing test bug: test__lp_linear_func_galaxy_dict_from was hidden by the old structural __dict__ eq on LightProfileLinear which made all parameter-less instances compare equal. New token-based eq exposed off-by-one variable references at lines 112-113.
• Pickling caveat (future work): Unpickled LightProfileLinear instances retain their pickled token. Acceptable for JIT usage; flag if staleness surfaces on long-lived pickles.

This unblocks the 11 remaining pytree variant prompts that ship models containing linear light profiles.