fix: latent/JAX release failures

[Jammy2211]

Overview

Fix the active latent/JAX release failures that still reproduce after the latent-class refactor merges. The work starts at the PyAutoFit-level JAX visualization dispatch assertion, then validates the autogalaxy workspace/test failures that depend on the same visualization and latent paths.

Plan

• Restore or replace the JAX visualization dispatch contract in a way that matches the current latent-refactor architecture.
• Verify the PyAutoFit dispatch assertion first because it is the smallest direct failure.
• Re-run the autogalaxy results and visualization failures to separate library-level regressions from workspace-script expectations.
• Patch workspace scripts only if the library contract is correct and the scripts are stale.
• Re-run the targeted Group A release failures, then the relevant unit/smoke checks.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• autogalaxy_workspace
• autogalaxy_workspace_test

Work Classification

Both: library-first, with workspace follow-up if the scripts need migration.

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	feature/latent-old-api-cleanup	dirty: autofit/example/analysis.py
./autogalaxy_workspace	main	clean
./autogalaxy_workspace_test	main	clean

Warning: the canonical PyAutoFit checkout has unregistered local work on feature/latent-old-api-cleanup. Development should use a task worktree from the current main commit so that local cleanup work is preserved.

Suggested branch: feature/latent-jax-release-failures
Worktree root: ~/Code/PyAutoLabs-wt/latent-jax-release-failures/

Implementation Steps

1. In PyAutoFit, inspect autofit/non_linear/analysis/analysis.py, autofit/non_linear/analysis/latent.py, and autofit/non_linear/fitness.py to confirm the intended post-refactor visualization dispatch path.
2. Fix Analysis.fit_for_visualization(...) or the workspace assertion contract so autofit_workspace_test/scripts/jax_assertions/fitness_dispatch.py passes without reintroducing stale state.
3. Run the PyAutoFit targeted tests covering analysis/latent/fitness dispatch.
4. Re-run these scripts against the task worktree environment:
   • autofit_workspace_test/scripts/jax_assertions/fitness_dispatch.py
   • autogalaxy_workspace/scripts/guides/results/start_here.py
   • autogalaxy_workspace_test/scripts/ellipse/modeling_visualization_jit.py
   • autogalaxy_workspace_test/scripts/ellipse/visualization_jax.py
   • autogalaxy_workspace_test/scripts/interferometer/modeling_visualization_jit.py
5. If autogalaxy workspace/test failures remain and the PyAutoFit contract is correct, route a workspace follow-up to migrate stale scripts or fix workspace data/config issues.

Key Files

• PyAutoFit/autofit/non_linear/analysis/analysis.py — JAX visualization dispatch entry point.
• PyAutoFit/autofit/non_linear/analysis/latent.py — refactored latent engine; may be involved in the empty-stack failure.
• PyAutoFit/autofit/non_linear/fitness.py — precompile/dispatch behaviour for JAX searches.
• autofit_workspace_test/scripts/jax_assertions/fitness_dispatch.py — smallest direct failing assertion.
• autogalaxy_workspace_test/scripts/*/modeling_visualization_jit.py and visualization_jax.py — downstream integration checks.

Original Prompt

Click to expand starting prompt

ok fix the first group

Context from adjusted report:

• autofit_workspace_test/scripts/jax_assertions/fitness_dispatch.py fails with Analysis has no attribute _jitted_fit_from.
• autogalaxy_workspace/scripts/guides/results/start_here.py fails with invalid/zero noise-map values.
• autogalaxy_workspace_test/scripts/ellipse/modeling_visualization_jit.py fails with expected jax.Array, got numpy.float64.
• autogalaxy_workspace_test/scripts/ellipse/visualization_jax.py fails because fit_ellipse.png is not produced.
• autogalaxy_workspace_test/scripts/interferometer/modeling_visualization_jit.py fails with compute_latent_samples empty stack.

[Jammy2211]

First Failure Group Fixed — 2026-06-09

PRs

• Library: fix: skip latent computation without keys #1317
• Workspace: test: update JAX visualization dispatch assertions autofit_workspace_test#33
• Workspace: fix: reload saved result dataset with unchecked noise map autogalaxy_workspace#112
• Workspace: test: align JAX visualization checks with current contract autogalaxy_workspace_test#65

Status

All first-group targeted reruns passed locally, and all PR CI checks are green. The PRs are labelled pending-release and are ready for library-first merge if approved.

Summary

The fixes cover the no-latent-key JAX crash, stale JAX visualization assertions, the saved-results guide noise-map reload, and missing ellipse JAX release-build environment overrides.

[Jammy2211]

Shipped — 2026-06-09

PRs

• fix: skip latent computation without keys #1317
• test: update JAX visualization dispatch assertions autofit_workspace_test#33
• fix: reload saved result dataset with unchecked noise map autogalaxy_workspace#112
• test: align JAX visualization checks with current contract autogalaxy_workspace_test#65

Summary

Merged the first grouped release-failure fixes after the latent refactor. The fix set covers the no-latent-key PyAutoFit crash, stale JAX visualization assertions, the Autogalaxy results guide noise-map reload, and missing ellipse JAX release-build environment overrides.

Verification

Targeted release reruns, PyAutoFit full suite, PyAutoBuild section checks, and all PR CI passed before merge. Kaplinghat branches were not touched.