feat: pytree-register FitImaging for jax.jit-returnable fits (Path A) #444

Affected Repositories

• PyAutoLens (primary — CLAUDE.md doc being lifted)
• PyAutoArray — AbstractNDArray pytree auto-registration restored
• PyAutoGalaxy — Galaxy, Redshift, light & mass profile pytree registration
• PyAutoFit — may need helper extension in autofit.jax.pytrees
• autolens_workspace_test — new mge_pytree.py integration test

Work Classification

Both (library-heavy, plus a workspace integration test).

Branch Survey

Suggested branch: feature/fit-imaging-pytree
Worktree root: ~/Code/PyAutoLabs-wt/fit-imaging-pytree/ (created later by /start_library)

Historical context — short-cut from git history

Pytree auto-registration already existed in PyAutoArray at tag 2025.12.15.2:

• autoarray/abstract_ndarray.py — AbstractNDArray.__init__ called register_pytree_node(type(self), instance_flatten, instance_unflatten) on first instantiation. Generic __dict__-walk with __no_flatten__ exclusion tuple. Reconstruction via __new__ + setattr — side-effect-free.
• Removed during the Nov 2025 xp refactor (6224b926 jax imports deferred, 30c6cbe6 follow-up, merged via a4762d7b). Motivation: defer JAX imports for non-JAX users.
• autoarray/operators/over_sampling/over_sampler.py still has a manual tree_flatten / tree_unflatten pair — mixed-mode precedent in current code.

PyAutoLens only ever pytree-registered point-solver classes (Step, ShapeSolver, NullTracer — 492a2e76b restore trees). PyAutoGalaxy has no prior pytree work.

Implementation Steps

1. PyAutoArray — restore AbstractNDArray pytree auto-registration, re-derived to match current code style. Gate behind use_jax=True:

   • Add a process-level flag set by Analysis.__init__ when use_jax=True (or reuse the existing autoconf.jax_wrapper state if that's what the registration should key off).
   • Re-implement instance_flatten / instance_unflatten in autoarray/abstract_ndarray.py matching the current codebase's typing and import conventions.
   • __no_flatten__ per-subclass tuple lists static attrs (masks, shapes, pixel_scales, config).
   • Add a targeted test in PyAutoArray/test_autoarray/jax/ following the three-step pattern from autolens_workspace_test/scripts/hessian_jax.py. One test is enough for the feasibility gate.

2. autolens_workspace_test — seed the driver script: copy scripts/jax_likelihood_functions/imaging/mge.py → mge_pytree.py, keep the same dataset / model / analysis, but replace fitness._vmap(parameters) with a jax.jit(analysis.fit_from) call on model.instance_from_prior_medians(). Assert the returned FitImaging's leaves are jax.Array, not NumPy.

3. Iterate on the traceback. Expected blockers (in likely order):

   • Grid2D, Mask2D — covered by step 1's auto-registration.
   • Galaxy, Redshift, light-profile dataclasses (Gaussian, Sersic, Exponential, etc.), mass-profile dataclasses (NFWSph, ExternalShear, PowerLaw), BasisLinear / GaussianGradient → register via autofit.jax.pytrees._build_instance_pytree_funcs pattern in PyAutoGalaxy.
   • Tracer, Galaxies (nested List[List[Galaxy]]) — confirm list recursion handles it, else wrap.
   • FitImaging itself — register last, children = the jax.Array-bearing attributes (model image, residual map, chi-squared map, etc.), aux = everything else.

4. Regression guard. PyAutoArray, PyAutoFit, PyAutoGalaxy, PyAutoLens full test suites must stay green. The gated registration must no-op when use_jax=False so non-JAX users see zero change.

5. Deliverable check. mge_pytree.py produces a JIT-wrapped FitImaging with jax.Array leaves; existing mge.py still asserts its expected likelihood; test suites pass.

Scope boundary

• No pixelization / inversion / Mapper / NNLS work (defer — mge.py's path doesn't touch them).
• No use_jax_for_visualization behaviour change (separate follow-up flips the dispatch to jax.jit).
• No edits to visualization_jax.py.
• Do NOT restore the "xp is np" guard removals — those are for a later task once this PoC succeeds.

Key Files

• PyAutoArray/autoarray/abstract_ndarray.py — restore pytree auto-registration, gated
• PyAutoArray/test_autoarray/jax/test_pytrees.py (new or existing) — PoC test
• PyAutoGalaxy/autogalaxy/galaxy/galaxy.py — register Galaxy via _build_instance_pytree_funcs
• PyAutoGalaxy/autogalaxy/profiles/ — register light & mass profile dataclasses
• PyAutoGalaxy/autogalaxy/cosmology/ — register Redshift if needed
• PyAutoLens/autolens/lens/tracer.py — register Tracer / Galaxies
• PyAutoLens/autolens/imaging/fit_imaging.py — register FitImaging (last)
• PyAutoGalaxy/autogalaxy/imaging/fit_imaging.py — base class may also need registration
• autolens_workspace_test/scripts/jax_likelihood_functions/imaging/mge_pytree.py — new
• PyAutoLens/CLAUDE.md — update the "autoarray types cannot be returned" restriction paragraph once PoC is green

Can we register @PyAutoLens/autolens/imaging/fit_imaging.py FitImaging (and the autoarray /
autogalaxy types it transitively contains) as JAX pytrees, so that a function that returns
FitImaging can be wrapped in jax.jit?

Why this matters

In the JAX-visualization pilot (#1227 on PyAutoFit, shipped via use_jax_for_visualization)
we intentionally took the "Path C" route: fit_from runs on the eager JAX path
(use_jax=True makes self._xp = jnp, operations run eagerly under JAX) and the
matplotlib plotters materialize arrays to NumPy at the boundary. No jax.jit wrapping.

The endgame — "Path A" — is full jax.jit wrapping of analysis.fit_from, so visualization
gets the same compile-time speedup the likelihood function does. This requires the
FitImaging return type (and every autoarray / autogalaxy / autolens type reachable from
it) to be a JAX pytree. PyAutoLens's CLAUDE.md currently documents the reverse:

Autoarray types (Array2D, ArrayIrregular, VectorYX2DIrregular, etc.) are not
registered as JAX pytrees. They can be constructed inside a JIT trace, but cannot
be returned as the output of a jax.jit-compiled function.

This task is the feasibility study for lifting that restriction.

What's already in place

Recent autofit work added autofit.jax.pytrees (see @PyAutoFit/autofit/jax/pytrees.py)
which registers Model, Collection, ModelInstance, and user classes found by walking
a Model tree. The pattern:

• children = dynamic (traced) arrays / sub-pytrees
• aux = static Python objects (concrete constants, class references, etc.)
• flatten / unflatten lift/restore the instance across the JIT boundary

The same pattern needs to extend down through the autoarray / autogalaxy layers.

Assess, don't implement (yet)

The deliverable is a feasibility assessment, not a fully-registered FitImaging. Produce
the assessment as a markdown document (e.g. in this prompt file's issued/ location, or
as a draft PR) answering:

1. Type inventory. Walk FitImaging (@PyAutoLens/autolens/imaging/fit_imaging.py) and
   list every distinct class reachable from a populated instance. For each class:

   • Where it's defined (autoarray / autogalaxy / autolens)
   • Whether it currently carries tree_flatten / tree_unflatten methods
   • Whether its constructor is compatible with pytree unflatten (takes the children as
     positional / keyword args) or whether a _build_*_pytree_funcs helper (like
     autofit.jax.pytrees._build_instance_pytree_funcs) would be needed

2. Dynamic vs static classification. For each class, which attributes are dynamic
   (JAX arrays — get traced) vs static (numpy arrays that mustn't change, masks, shapes,
   pixel scales, redshifts, config objects)? The CLAUDE.md rule for autoarray is:
   grid.array[:, 0] is traced, grid.mask is static. That rule needs to be restated
   per type here.

3. Structural blockers. Are there types that cannot be registered cleanly? Examples
   to look for:

   • Classes whose __init__ does non-trivial work that can't be replayed during
     unflatten (e.g. reading config, triggering a numba compilation)
   • Classes that hold live references to an Inversion whose solver state isn't
     pickleable / is stateful under JAX
   • Tracer / Galaxies — these are List[List[Galaxy]] nested containers; are they
     already pytree-friendly via Python list recursion, or do they need explicit
     registration?
   • Mapper / LinearEqn / NNLS solver state inside Inversion — likely the hardest

4. Interaction with the existing "xp is np" guard pattern. PyAutoLens has an established
   convention (see @PyAutoGalaxy/autogalaxy/operate/lens_calc.py): functions that return
   autoarray wrappers guard with if xp is np: return Array2D(...) else return raw
   jax.Array. Registering Array2D as a pytree means these guards may become
   unnecessary. Is removing them desirable, or should they stay for the non-JIT path?

5. Proof-of-concept. Register one concrete type (suggest Array2D — it's the
   simplest leaf) and demonstrate that a toy function def f(x): return Array2D(values=x*2, mask=static_mask) runs under jax.jit and returns an Array2D backed by
   jax.Array. This both validates the approach and surfaces any wrinkles that the static
   analysis missed.

6. Estimate. Given the inventory and the proof-of-concept, roughly how many types need
   registration to make FitImaging a pytree? Which ones are cheap (static dataclass-ish),
   which are hairy (inversion, mapper)?

Scope boundary

• Do not change use_jax_for_visualization behaviour. That flag currently dispatches
  to the eager-JAX path; once this task produces a green light + registration PR, a
  follow-up task will flip the dispatch to jax.jit.
• Do not start mass pytree-registration in this task. One-type PoC only.
• Do not register FitImaging itself yet — that's the goal after every type it
  reaches is registered.

Starting points

• @PyAutoLens/autolens/imaging/fit_imaging.py — the FitImaging subclass
• @PyAutoGalaxy/autogalaxy/imaging/fit_imaging.py — base FitImaging
• @PyAutoArray/autoarray/structures/arrays/uniform_2d.py — Array2D (suggested PoC type)
• @PyAutoFit/autofit/jax/pytrees.py — the existing autofit pytree machinery; mirror its
  structure when proposing new registrations
• @PyAutoLens/CLAUDE.md — the "xp is np" guard rule this task's outcome will revise

Deliverables

1. Type inventory + classification table (markdown)
2. List of structural blockers with suggested workarounds
3. Working PoC registering Array2D with a minimal test in
   @PyAutoArray/test_autoarray/jax/ following the three-step pattern in
   autolens_workspace_test/scripts/hessian_jax.py
4. Effort estimate + recommended ordering for the follow-up implementation task