refactor: migrate imaging/mge.py JIT closure to pytree parameter inputs

[Jammy2211]

Overview

The flagship JAX profiling script jax_profiling/imaging/mge.py currently JITs a closure that takes a flat 1D parameter vector and unpacks it via model.instance_from_vector inside the trace. This obscures parameter identity under jax.jit, forces vmap callers to stack (batch, N) arrays, and does not match the pytree-native direction the rest of the codebase has moved toward. This task migrates the script to pytree inputs so JIT and vmap flatten/unflatten parameters automatically and preserve their names.

Plan

• Dig through ~12-month-old git history for prior pytree implementations in PyAutoFit, PyAutoArray and PyAutoGalaxy; use them as templates rather than designing from scratch.
• Re-enable / extend pytree registration on af.ModelInstance (and supporting classes in autogalaxy / autoarray as needed) so it can cross a jax.jit boundary directly.
• Replace the flat jnp_params JIT input in mge.py with a structured pytree.
• Update mapping_matrix_from_params, blurred_mm_from_params, the full-pipeline JIT and the vmap path to consume the pytree directly — removing model.instance_from_vector from inside the trace.
• Keep profiling outputs (per-step timings, JSON summary, bar chart) format-compatible; timings should match within noise.
• Refresh the module docstring to document pytree inputs as the recommended pattern for new profiling scripts.
• Follow-up (separate work): propagate the same pattern to pixelization.py and delaunay.py.

Detailed implementation plan

Affected Repositories

• PyAutoLabs/autolens_workspace_developer (primary — originating request)
• rhayes777/PyAutoFit (pytree registration on PriorModel / ModelInstance)
• Jammy2211/PyAutoArray (potentially — extend existing uniform.py / triangle pytree wrapping)
• Jammy2211/PyAutoGalaxy (potentially — Galaxy, Tracer, light/mass profiles)

Work Classification

Library + Workspace. /start_library first; autolens_workspace_developer script changes follow once library-side registration is in place.

Historical pytree references

Repo	Commit / PR	Content
PyAutoFit	autofit/mapper/prior_model/prior_model.py:281,289	tree_flatten / tree_unflatten on PriorModel already present; register_pytree_node call at line 211 is currently commented out — check git blame for the reason before re-enabling
PyAutoFit	PR #1037 feature/jax_pytree, PR #1103 feature/graphical_pytrees	AnalysisFactor, FactorGraphModel pytree methods, prior-id children
PyAutoArray	9a8d17fe "tree flatten point and circle" (Sep 2024)	autoarray/structures/triangles/shape.py
PyAutoArray	e6419042 + PR #137 feature/jax_tracer (Oct 2024)	autoarray/operators/over_sampling/uniform.py — still present today
PyAutoGalaxy	feature/jax_mge, feature/jax_tracer branches	precursor profile/tracer pytree work

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_developer	main	dirty (unrelated fits / json / png results + untracked probe scripts; snapshotted to z_staging/imaging-mge-pytree-migration-dirty-snapshot/; do not touch)
./PyAutoFit	main	(check at worktree creation)
./PyAutoArray	main	(check at worktree creation)
./PyAutoGalaxy	main	(check at worktree creation)

Suggested branch: feature/imaging-mge-pytree-migration (used consistently across all affected repos)
Worktree root: ~/Code/PyAutoLabs-wt/imaging-mge-pytree-migration/ (created later by /start_library)

Implementation Steps

1. git show / git diff the historical commits listed above to understand the prior pytree approach. Read autofit/mapper/prior_model/prior_model.py:205-295 and investigate why register_pytree_node was commented out.
2. Re-enable register_pytree_node on PriorModel (and ModelInstance if needed) in PyAutoFit. If the reason for the original comment-out still applies, work around it (e.g. selective leaf registration) rather than blindly uncommenting.
3. Extend pytree registration to Galaxy, Tracer, and the lp_linear.MultiGaussian / mp.Isothermal / mp.ExternalShear profile classes used in mge.py — only as far as needed to let a ModelInstance cross a JIT boundary and be rebuilt into a Tracer inside the trace.
4. Add unit tests in PyAutoFit / PyAutoGalaxy covering flatten → unflatten round-trip and jax.jit / jax.vmap over a simple ModelInstance.
5. In autolens_workspace_developer/jax_profiling/imaging/mge.py:
   • After instance = model.instance_from_vector(vector=param_vector) (setup, not JIT), pass instance (pytree-registered) directly as the JIT input.
   • Rewrite mapping_matrix_from_params (mge.py:336): accept the pytree, build Tracer via al.Tracer(galaxies=list(instance.galaxies)) — no instance_from_vector in the trace.
   • Rewrite blurred_mm_from_params (mge.py:373): same change.
   • Replace the full-pipeline fitness.call(jnp_params) profile block (mge.py:554-563) with a pytree-accepting closure defined inside the script. Do NOT modify Fitness.call itself (explicit out-of-scope per original prompt).
   • vmap section (mge.py:572-603): build batched pytree via jax.tree_util.tree_map(lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)), instance), vmap the new closure, keep the correctness assertion.
   • Static memory analysis block (mge.py:609-621): point batched_call at the pytree closure.
   • Update the top-of-file docstring and inline comments to document pytree inputs as the recommended pattern.
6. Verify: cd jax_profiling/imaging && NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python mge.py. Correctness asserts must pass; timings within noise of prior JSON results.

Key Files

• PyAutoFit/autofit/mapper/prior_model/prior_model.py — re-enable pytree registration
• PyAutoGalaxy/autogalaxy/galaxy/galaxy.py, .../profiles/** — extend registration as needed
• PyAutoLens/autolens/lens/tracer.py — extend registration as needed
• autolens_workspace_developer/jax_profiling/imaging/mge.py — the migration

Out of scope (explicit)

• PyAutoFit/autofit/non_linear/fitness.py (Fitness.call)
• PyAutoFit model.instance_from_vector (production path)
• jax_profiling/imaging/pixelization.py — follow-up
• jax_profiling/imaging/delaunay.py — follow-up
• Any dirty/untracked files in autolens_workspace_developer (snapshotted to z_staging/imaging-mge-pytree-migration-dirty-snapshot/; left alone)

Original Prompt

Click to expand starting prompt

# Imaging MGE JAX JIT Profiling — Migrate to Pytree Inputs

## Context

`autolens_workspace_developer/jax_profiling/imaging/mge.py` is the
flagship JAX JIT profiling script for the MGE source model. It
builds `Fitness.call` around a 1D parameter vector
(`jnp_params`, shape `(N,)`) and JITs that closure.

Passing parameters as a flat 1D vector has two long-term downsides:

1. **`model.instance_from_vector` is the JIT boundary**. That call
   has to unpack the flat vector into every named structured
   parameter (centres, ell_comps, intensities, sigmas, redshifts,
   …) inside the trace. For a model with many parameters this is
   pure Python shape-juggling that runs once per compile but
   obscures the parameter identity under `jax.jit`.
2. **Batching with `vmap`** requires the caller to stack parameter
   vectors into a 2D `(batch, N)` array. The structured parameter
   names are lost along the way — any diagnostic that wants to
   plot likelihood vs `einstein_radius` has to know the positional
   index of `einstein_radius` in the flat vector, which is
   fragile.

The rest of the codebase (tracers, galaxies, profiles) has moved
toward **pytree-native** inputs: nested dicts / `ModelInstance`
objects registered as JAX pytrees, so `jax.jit` flattens and
unflattens them automatically, and `vmap` batches over the pytree
leaves without the caller having to reshape anything.

## Task

Update `autolens_workspace_developer/jax_profiling/imaging/mge.py`
so that:

1. The JIT'd closure takes a **pytree** of parameters (e.g. a
   nested dict keyed by galaxy name → profile name → parameter
   name, or the `af.ModelInstance` object itself if it's now
   pytree-registered) rather than a flat 1D array.
2. `jax.jit` and `jax.vmap` are applied to that pytree-accepting
   closure directly — no flat-vector shim.
3. The profiling output (compile time, steady-state runtime, vmap
   batch throughput) should be unchanged within noise — this is a
   clarity / API-ergonomics change, not a performance change.
4. Update the inline commentary to point at the pytree inputs as
   the recommended pattern for new profiling scripts.

Once `mge.py` is updated, treat it as the new reference style and
propagate the same pattern to `pixelization.py` and `delaunay.py`
as a follow-up.

## Why this matters

- Downstream users reading the script will more easily see *which*
  parameter is which, because the pytree preserves names.
- `jax.jit` compile caches are keyed on the abstract pytree
  structure; this change makes the cache key human-readable.
- Eventually we will want to build user-facing `jax.jit`'d
  likelihood functions that accept a `ModelInstance` directly —
  the profiling script should demonstrate that end-state, not an
  older flat-vector style.

## Risks / blockers

- If `af.ModelInstance` is not yet registered as a JAX pytree,
  this change either (a) needs that registration to happen first,
  or (b) uses a plain nested dict as an interim representation.
  Check `autofit`'s pytree registration status before starting;
  file a blocker issue if registration is missing.
- `model.instance_from_vector(..., xp=jnp)` is currently relied on
  by several other profiling scripts and by the production
  `Fitness.call`. This task *only* changes the profiling script,
  not the production path — do not touch `Fitness` or `model.instance_from_vector`.

## Pytree Registration Old

Older versions (12ish months ago) of PyAutoFit, PyAutoArray and PyAutoGalaxy had working Pytree circulations.
Dig through their git history first and see if you can find source code there to help you as a reference.

[Jammy2211]

Library PRs Created — 2026-04-16

PRs

• feat: add lazy JAX pytree registration helpers to autoconf PyAutoConf#93
• feat: add autofit.jax opt-in pytree registration API PyAutoFit#1220

Status

Workspace changes needed. Next: /start_workspace to ship the mge.py migration in autolens_workspace_developer.

API Changes

Purely additive — no removals or signature changes.

• New autoconf.jax_wrapper.register_pytree_node_class(cls) / register_pytree_node(nodetype, flatten_func, unflatten_func) — lazy, no-op if JAX missing.
• New autofit.jax sub-package exposing enable_pytrees() and register_model(model). The latter walks a Model / Collection and partitions each Model.cls's constructor args: constants (from direct_instance_tuples) → JAX aux_data so they stay concrete under JIT; priors → children. This keeps code like sorted(galaxies, key=lambda g: g.redshift) working inside jax.jit.

Workspace Usage

autolens_workspace_developer/jax_profiling/imaging/mge.py now runs its full vmap(jit(full_pipeline_from_params)) pipeline on a pytree ModelInstance instead of a flat 1D vector — no more instance_from_vector inside the trace. Fitness.call is untouched.

[Jammy2211]

Shipped — 2026-04-16

PRs

• feat: migrate jax_profiling/imaging/mge.py to pytree inputs #11 (workspace)

Upstream PRs

• feat: add lazy JAX pytree registration helpers to autoconf PyAutoConf#93
• feat: add autofit.jax opt-in pytree registration API PyAutoFit#1220

Summary

Migrated jax_profiling/imaging/mge.py from a flat 1D parameter vector to pytree ModelInstance inputs. Full-pipeline JIT + vmap now consume the pytree directly via analysis.log_likelihood_function(instance=params_tree) with use_jax=True. Fitness.call and model.instance_from_vector are untouched — the production sampler path is unaffected.

The library side adds autofit.jax with enable_pytrees() + register_model(model). register_model partitions each Model.cls constructor's arguments: constants (from direct_instance_tuples) go into JAX aux_data so they stay concrete under JIT, priors become children. This keeps code like sorted(galaxies, key=lambda g: g.redshift) working inside jax.jit.

Session Notes

• Constants-as-aux_data was the critical insight. The first cut stored every constructor argument as a child. Under JIT, Galaxy.redshift=0.5 became a tracer and Tracer.galaxies_ascending_redshift blew up on sorted(..., key=lambda g: g.redshift) with TracerBoolConversionError. Routing constants from model.direct_instance_tuples into aux_data fixed it.
• use_jax=True on the Analysis, not FitImaging directly. An earlier attempt used al.FitImaging(dataset, tracer, settings) which ran xp=np internally and raised TracerArrayConversionError inside autogalaxy/convert.py (np.where on a JAX bool). Switching to al.AnalysisImaging(dataset, use_jax=True) + analysis.log_likelihood_function(...) threads xp=jnp through the whole pipeline.
• Lazy, opt-in registration preserves the Nov 2025 cleanup. The eager register_pytree_node_class calls that forced a jax.tree_util import on every import autofit are intentionally not restored. enable_pytrees() defers the JAX import to call time and no-ops if JAX is missing; ValueErrors from duplicate registration are tolerated silently.
• Future work (out of scope here): propagate the same pytree pattern to jax_profiling/imaging/pixelization.py and delaunay.py. Called out in the original prompt as follow-up.

[Jammy2211]

Library PR Created — 2026-04-17

PR

• fix(jax): route concrete kwarg attrs to aux_data in pytree flatten PyAutoFit#1221

Status

Workspace changes needed. Next: /ship_workspace (workspace script migration is already implemented; this PR ships it).

Summary

Fixes two bugs in autofit.jax._build_instance_pytree_funcs that surfaced when migrating jax_profiling/imaging/pixelization.py to pass a ModelInstance pytree into JIT/vmap directly:

1. Concrete non-Prior kwargs (e.g. Galaxy(pixelization=<Pixelization>)) were routed into children instead of aux_data, so under jax.jit the pixelization became a tracer and isinstance(x, aa.Pixelization) silently returned False — leaving fit.inversion None.
2. **kwargs classes like Galaxy produce instances with different attribute sets across models (lens: bulge/mass/shear; source: pixelization). JAX allows one flatten per class, so the first model's captured name list raised AttributeError when flattening the other instance.

Fix

Shared _CLASS_FIELD_CLASSIFIERS[cls] dict accumulates classifications across every register_model walk. flatten now iterates vars(instance) at call time and defaults unknown attrs to aux_data.

[Jammy2211]

Shipped (pixelization) — 2026-04-17

PRs

• Workspace: refactor(jax_profiling): migrate imaging/pixelization.py to pytree inputs #12
• Library: fix(jax): route concrete kwarg attrs to aux_data in pytree flatten PyAutoFit#1221

Summary

jax_profiling/imaging/pixelization.py now passes a ModelInstance pytree directly into jit/vmap via autofit.jax.register_model, matching the mge.py pattern shipped in #11. Both correctness assertions pass end-to-end (step-by-step log_evidence vs reference rtol=1e-4; vmap batch=3 vs single JIT rtol=1e-4).

Session Notes

• Scope expanded mid-session from workspace-only to library + workspace. The first smoke-test run surfaced a bug in autofit.jax._build_instance_pytree_funcs: concrete non-Prior kwargs (e.g. Galaxy(pixelization=<Pixelization>)) were classified using direct_argument_names / direct_instance_tuples, which don't see attributes stored directly on model.__dict__. They ended up in children and became JAX tracers under JIT, breaking isinstance(x, aa.Pixelization) inside FitImaging so fit.inversion silently resolved to None.
• First library fix (using model.items() for classification) was still insufficient: Galaxy.__init__(self, redshift, **kwargs) produces instances with different attribute sets per model (lens: bulge/mass/shear; source: pixelization) while sharing one cls. JAX registers flatten per-class, so the captured name list from whichever model was seen first raised AttributeError: 'Galaxy' object has no attribute 'bulge' when flattening the other instance.
• Second library fix (shipped): shared _CLASS_FIELD_CLASSIFIERS[cls] dict accumulates classifications across every register_model walk via setdefault; flatten iterates vars(instance) at call time and defaults unknown attrs to aux_data (safer than tracing an unknown object). Regression test test_register_model_keeps_kwarg_constants_static covers the concrete-kwarg + **kwargs pattern.
• Left alone: per-step raw-array JIT closures (ray_trace, inversion linear algebra) — no params → out of scope. FitImaging(..., xp=jnp) kept inside blurred_mm_from_params because that closure extracts operated_mapping_matrix which AnalysisImaging.log_likelihood_function does not expose.

[Jammy2211]

Shipped (follow-up) — 2026-04-17

PRs

• Migrate mge_gradients.py JAX profiling to pytree inputs #13

Summary

Migrated jax_profiling/imaging/mge_gradients.py to pass a ModelInstance pytree through jit/value_and_grad instead of a flat 1D parameter vector, matching the pattern shipped in #11 (mge.py) and #12 (pixelization.py). The Fitness wrapper and Model.instance_from_vector no longer appear inside the JIT trace — gradient profiling now reflects how a real pytree-native sampler would run.

Session Notes

• PART C full-pipeline gradient swapped from Fitness.call to AnalysisImaging(use_jax=True).log_likelihood_function(instance=params_tree), same substitution used in mge.py (feat: migrate jax_profiling/imaging/mge.py to pytree inputs #11). pixelization.py (refactor(jax_profiling): migrate imaging/pixelization.py to pytree inputs #12) kept FitImaging(..., xp=jnp) inside its per-step closure because it specifically needs operated_mapping_matrix; mge_gradients.py's per-step closures don't need inversion internals, so AnalysisImaging fits cleanly for the full-pipeline step only.
• test_grad helper now tree-flattens the gradient via jax.tree_util.tree_leaves(grad) + np.concatenate(...) for statistics — gradient shape is no longer a flat 1D vector.
• NNLS _diagnose_kappa inner _loss uses jax.value_and_grad(_loss)(params_tree) with the same tree-flatten pattern; all four kappa values (1e-3, 1e-2, 1e-1, 1.0) reported FULLY FINITE GRADIENTS.
• CLAUDE.md gotcha: script is run from the workspace root (not jax_profiling/imaging/) because dataset_path is relative to pwd. The CLAUDE.md "cd jax_profiling/imaging" hint is misleading for this specific script.