You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
PyAutoLens visualizers were updated in PR #443 (2026-04-19) to dispatch through analysis.fit_for_visualization(instance=instance) so the fit reused for plotting goes through autofit's cached jax.jit wrapper when use_jax_for_visualization=True. PyAutoGalaxy received matching pytree registration for FitImaging (PR #364, 2026-04-22) and FitInterferometer (PR #376), but the visualizer dispatch was never switched over for either. As a result use_jax_for_visualization=True on ag.AnalysisImaging / AnalysisInterferometer is currently a silent no-op for visualization. This task swaps the three call sites.
Plan
Three fit_from(instance=...) calls in PyAutoGalaxy visualizer files become fit_for_visualization(instance=...).
The change is safe for the NumPy default path because fit_for_visualization falls back to fit_from whenever use_jax_for_visualization=False.
Existing imaging JAX visualization smoke test now actually exercises the jit-cached path instead of silently no-opping.
No new pytree registration in this task — both imaging and interferometer pytrees are already shipped.
Detailed implementation plan
Affected Repositories
PyAutoGalaxy (primary)
autogalaxy_workspace_test (verification only — no edits)
Work Classification
Library
Branch Survey
Repository
Current Branch
Dirty?
./PyAutoGalaxy
main
clean
Suggested branch:feature/autogalaxy-viz-dispatch-swap Worktree root:~/Code/PyAutoLabs-wt/autogalaxy-viz-dispatch-swap/ (created later by /start_library)
Implementation Steps
autogalaxy/imaging/model/visualizer.py:79 — replace fit = analysis.fit_from(instance=instance) with fit = analysis.fit_for_visualization(instance=instance) (single-analysis visualize path).
autogalaxy/imaging/model/visualizer.py:176 — same replacement inside the visualize_combined list-comprehension that builds per-analysis fits for multi-analysis fits.
autogalaxy/interferometer/model/visualizer.py:81 — same replacement (single-analysis visualize path; no visualize_combined exists in the interferometer visualizer).
Run pytest test_autogalaxy/imaging and pytest test_autogalaxy/interferometer to confirm no regressions.
Run /smoke_test autogalaxy_workspace_test imaging/visualization.py imaging/visualization_jax.py interferometer/visualization.py — all three must pass. The visualization_jax.py script will now actually exercise the jit-cached path rather than silently no-opping.
PyAutoLens visualizers were updated to dispatch through analysis.fit_for_visualization(instance=instance)
in PR #443 (commit 761c1778e, 2026-04-19) so that — when use_jax_for_visualization=True is set on
the analysis — the fit reused for plotting goes through the cached jax.jit wrapper added to the
autofit base Analysis in PR #1228.
PyAutoGalaxy was later given matching JAX pytree registration for FitImaging + DatasetModel + Galaxies (PR #364, 2026-04-22) but the visualizer dispatch was never switched over. As a result use_jax_for_visualization=True on ag.AnalysisImaging is currently a no-op for visualization —
the pytree machinery is in place but the call sites still go through the eager analysis.fit_from(instance=instance) path.
What to change
@PyAutoGalaxy/autogalaxy/imaging/model/visualizer.py:79 — swap fit = analysis.fit_from(instance=instance) for fit = analysis.fit_for_visualization(instance=instance).
@PyAutoGalaxy/autogalaxy/imaging/model/visualizer.py:176 — same swap inside the visualize_combined path that builds per-analysis fits in a multi-analysis scenario.
@PyAutoGalaxy/autogalaxy/interferometer/model/visualizer.py:81 — same swap.
Note: AnalysisInterferometer pytree registration shipped in PyAutoGalaxy PR feat: register pytrees for AnalysisInterferometer #376
(see complete.md 2026-04-? entry — _register_fit_interferometer_pytrees is now
active in autogalaxy/interferometer/model/analysis.py:165-184), so dispatching
through fit_for_visualization is now safe for interferometer. Originally this
bullet was deferred — it was added to scope on 2026-05-08 once interferometer
pytree registration landed.
fit_for_visualization is defined on the autofit base Analysis. It dispatches to a jax.jit-cached wrapper when use_jax_for_visualization=True and falls back to plain fit_from
otherwise — so this change is safe for the NumPy default path.
autogalaxy_workspace_test/scripts/imaging/visualization_jax.py should still pass (uses use_jax_for_visualization=True) and will now actually exercise the jit-cached path rather
than silently no-op.
autogalaxy_workspace_test/scripts/imaging/visualization.py should still pass (NumPy path — fit_for_visualization falls back to fit_from when the flag is off).
Run both via /smoke_test.
For interferometer the equivalent workspace_test scripts don't exist yet — coverage is added
by autogalaxy_workspace_test/jax_viz_dataset_coverage.md (Phase 1C of the JAX
visualization roadmap). Verify the dispatch swap doesn't break existing autogalaxy_workspace_test/scripts/interferometer/ (NumPy path — should be unchanged).
Out of scope
Production workspace adoption (autogalaxy_workspace scripts opting into use_jax_for_visualization=True) — defer until Path A from issued/fit_imaging_pytree.md lands.
AnalysisEllipse / AnalysisQuantity — these have no pytree registration today;
covered by autogalaxy/fit_pytree_registration_other_datasets.md (which is now
scoped to ellipse + quantity only — interferometer pytree registration shipped in
PR feat: register pytrees for AnalysisInterferometer #376).
Background
Original feature: complete.md entries jax-visualization and mge-jit-visualization
(both 2026-04-19). Autogalaxy imaging pytree registration: complete.md entry for the
imaging port (2026-04-22).
Overview
PyAutoLens visualizers were updated in PR #443 (2026-04-19) to dispatch through
analysis.fit_for_visualization(instance=instance)so the fit reused for plotting goes through autofit's cachedjax.jitwrapper whenuse_jax_for_visualization=True. PyAutoGalaxy received matching pytree registration forFitImaging(PR #364, 2026-04-22) andFitInterferometer(PR #376), but the visualizer dispatch was never switched over for either. As a resultuse_jax_for_visualization=Trueonag.AnalysisImaging/AnalysisInterferometeris currently a silent no-op for visualization. This task swaps the three call sites.Plan
fit_from(instance=...)calls in PyAutoGalaxy visualizer files becomefit_for_visualization(instance=...).fit_for_visualizationfalls back tofit_fromwheneveruse_jax_for_visualization=False.Detailed implementation plan
Affected Repositories
Work Classification
Library
Branch Survey
Suggested branch:
feature/autogalaxy-viz-dispatch-swapWorktree root:
~/Code/PyAutoLabs-wt/autogalaxy-viz-dispatch-swap/(created later by/start_library)Implementation Steps
autogalaxy/imaging/model/visualizer.py:79— replacefit = analysis.fit_from(instance=instance)withfit = analysis.fit_for_visualization(instance=instance)(single-analysisvisualizepath).autogalaxy/imaging/model/visualizer.py:176— same replacement inside thevisualize_combinedlist-comprehension that builds per-analysis fits for multi-analysis fits.autogalaxy/interferometer/model/visualizer.py:81— same replacement (single-analysisvisualizepath; novisualize_combinedexists in the interferometer visualizer).pytest test_autogalaxy/imagingandpytest test_autogalaxy/interferometerto confirm no regressions./smoke_test autogalaxy_workspace_test imaging/visualization.py imaging/visualization_jax.py interferometer/visualization.py— all three must pass. Thevisualization_jax.pyscript will now actually exercise the jit-cached path rather than silently no-opping.Key Files
autogalaxy/imaging/model/visualizer.py— two-line change (lines 79, 176)autogalaxy/interferometer/model/visualizer.py— one-line change (line 81)Reference (PyAutoLens equivalents — pattern to mirror)
PyAutoLens/autolens/imaging/model/visualizer.py:97— single-analysis dispatchPyAutoLens/autolens/imaging/model/visualizer.py:239— multi-analysis combined dispatchPyAutoLens/autolens/interferometer/model/visualizer.py:96+:209Original Prompt
Click to expand starting prompt
PyAutoLens visualizers were updated to dispatch through
analysis.fit_for_visualization(instance=instance)in PR #443 (commit 761c1778e, 2026-04-19) so that — when
use_jax_for_visualization=Trueis set onthe analysis — the fit reused for plotting goes through the cached
jax.jitwrapper added to theautofit base
Analysisin PR #1228.PyAutoGalaxy was later given matching JAX pytree registration for
FitImaging+DatasetModel+Galaxies(PR #364, 2026-04-22) but the visualizer dispatch was never switched over. As a resultuse_jax_for_visualization=Trueonag.AnalysisImagingis currently a no-op for visualization —the pytree machinery is in place but the call sites still go through the eager
analysis.fit_from(instance=instance)path.What to change
@PyAutoGalaxy/autogalaxy/imaging/model/visualizer.py:79— swapfit = analysis.fit_from(instance=instance)forfit = analysis.fit_for_visualization(instance=instance).@PyAutoGalaxy/autogalaxy/imaging/model/visualizer.py:176— same swap inside thevisualize_combinedpath that builds per-analysis fits in a multi-analysis scenario.@PyAutoGalaxy/autogalaxy/interferometer/model/visualizer.py:81— same swap.Note:
AnalysisInterferometerpytree registration shipped in PyAutoGalaxy PR feat: register pytrees for AnalysisInterferometer #376(see
complete.md2026-04-? entry —_register_fit_interferometer_pytreesis nowactive in
autogalaxy/interferometer/model/analysis.py:165-184), so dispatchingthrough
fit_for_visualizationis now safe for interferometer. Originally thisbullet was deferred — it was added to scope on 2026-05-08 once interferometer
pytree registration landed.
fit_for_visualizationis defined on the autofit baseAnalysis. It dispatches to ajax.jit-cached wrapper whenuse_jax_for_visualization=Trueand falls back to plainfit_fromotherwise — so this change is safe for the NumPy default path.
Reference (PyAutoLens equivalent)
@PyAutoLens/autolens/imaging/model/visualizer.py:97— single-analysis dispatch@PyAutoLens/autolens/imaging/model/visualizer.py:239— multi-analysisvisualize_combineddispatchVerification
autogalaxy_workspace_test/scripts/imaging/visualization_jax.pyshould still pass (usesuse_jax_for_visualization=True) and will now actually exercise the jit-cached path ratherthan silently no-op.
autogalaxy_workspace_test/scripts/imaging/visualization.pyshould still pass (NumPy path —fit_for_visualizationfalls back tofit_fromwhen the flag is off)./smoke_test.by
autogalaxy_workspace_test/jax_viz_dataset_coverage.md(Phase 1C of the JAXvisualization roadmap). Verify the dispatch swap doesn't break existing
autogalaxy_workspace_test/scripts/interferometer/(NumPy path — should be unchanged).Out of scope
use_jax_for_visualization=True) — defer until Path A fromissued/fit_imaging_pytree.mdlands.AnalysisEllipse/AnalysisQuantity— these have no pytree registration today;covered by
autogalaxy/fit_pytree_registration_other_datasets.md(which is nowscoped to ellipse + quantity only — interferometer pytree registration shipped in
PR feat: register pytrees for AnalysisInterferometer #376).
Background
Original feature:
complete.mdentriesjax-visualizationandmge-jit-visualization(both 2026-04-19). Autogalaxy imaging pytree registration:
complete.mdentry for theimaging port (2026-04-22).