fix: skip _compute_latent_samples in PYAUTO_TEST_MODE

[Jammy2211]

Overview

After the latent refactor series shipped (z_features/complete/latent_refactor.md), the autolens SLaM smoke test times out at the 5-minute mark inside the first SLaM stage's post-fit latent processing. Root cause: PyAutoFit:autofit/non_linear/search/updater.py:226 (_compute_latent_samples) ignores PYAUTO_TEST_MODE. It always draws output.latent_draw_via_pdf_size samples (default 100) and runs analysis.compute_latent_samples(...) on each — and PyAutoGalaxy's AnalysisDataset.LATENT_BATCH_MODE = "jit" override means that's a Python loop of per-sample JIT dispatches because the einstein-radius latent's ZeroSolver isn't vmap-compatible.

Test mode by definition mocks the fit, so the latent values are never asserted on — short-circuiting _compute_latent_samples in test mode unblocks smoke testing immediately with no behavioural change to real runs.

Plan

• Add a skip_latents() helper in PyAutoConf:autoconf/test_mode.py (alongside existing skip_visualization(), skip_checks() etc) that composes is_test_mode() OR PYAUTO_SKIP_LATENTS=1.
• In PyAutoFit:autofit/non_linear/search/updater.py, short-circuit _compute_latent_samples to return None when skip_latents() is true.
• Add unit test covering both the test-mode trigger and the explicit PYAUTO_SKIP_LATENTS=1 trigger.
• Validate against the previously-timing-out SLaM smoke command.

Detailed implementation plan

Affected Repositories

• rhayes777/PyAutoFit (primary — updater.py short-circuit)
• rhayes777/PyAutoConf (test_mode.py adds skip_latents() helper)

Work Classification

Library.

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./PyAutoConf	main	clean

Suggested branch: feature/test-mode-skip-latents
Worktree root: ~/Code/PyAutoLabs-wt/test-mode-skip-latents/ (created later by /start_library)

Implementation Steps

1. PyAutoConf — add skip_latents() to autoconf/test_mode.py:

   def skip_latents():
       """Return True if latent variable computation should be skipped."""
       return is_test_mode() or os.environ.get("PYAUTO_SKIP_LATENTS", "0") == "1"

   Place after skip_checks() to match the file's existing layout.

2. PyAutoFit — short-circuit _compute_latent_samples in autofit/non_linear/search/updater.py:226:

   def _compute_latent_samples(self, ...):
       """Compute and persist latent variable samples if configured."""
       from autoconf.test_mode import skip_latents
       if skip_latents():
           return None
       if not ((during_analysis and conf.instance["output"]["latent_during_fit"])
               or (not during_analysis and conf.instance["output"]["latent_after_fit"])):
           return
       ...

   _profile_and_summarize (line 117 callsite) already accepts latent_samples=None, so no downstream changes.

3. PyAutoFit unit test in test_autofit/non_linear/search/ (confirm test layout first; if no test_updater.py exists, add one or use the appropriate adjacent file). Three cases:

   • With PYAUTO_TEST_MODE=1 set → _compute_latent_samples returns None without invoking analysis.compute_latent_samples.
   • With PYAUTO_SKIP_LATENTS=1 set (test mode off) → same behaviour.
   • Default (both unset) → existing path still runs (mock analysis.compute_latent_samples to return a sentinel; assert it's called).

   No JAX in the unit test.

4. Validate with the SLaM smoke command that previously timed out:

   PYAUTO_TEST_MODE=2 PYAUTO_SKIP_FIT_OUTPUT=1 PYAUTO_SKIP_VISUALIZATION=1 \
   PYAUTO_SKIP_CHECKS=1 PYAUTO_SMALL_DATASETS=1 PYAUTO_FAST_PLOTS=1 \
   python scripts/imaging/features/advanced/slam/start_here.py

   Should complete in seconds. Then a no-PYAUTO_TEST_MODE real-mode fit to confirm latent.csv still produced as before.

Key Files

• PyAutoConf/autoconf/test_mode.py — add skip_latents() helper
• PyAutoFit/autofit/non_linear/search/updater.py:226 (_compute_latent_samples) — short-circuit
• PyAutoFit/test_autofit/non_linear/search/ (test layout TBD) — unit test

Out of Scope

• Per-latent LATENT_BATCH_MODE (so fast latents vmap while einstein-radius stays jit-per-sample). That's the real-mode SLaM perf fix and lives in a separate follow-up at PyAutoPrompt/autogalaxy/latent_batch_mode_per_latent.md (TBA). This issue only unblocks smoke testing.
• No new workspace CLAUDE.md documentation: the existing fast-smoke env set covers test mode; PYAUTO_SKIP_LATENTS=1 is for power users and is discoverable via autoconf/test_mode.py.

Original Prompt

Click to expand starting prompt

Located at PyAutoPrompt/issued/test_mode_skip_latents.md after routing.

[Jammy2211]

Smoke Test Results — 2026-05-24

Worktree: feature/test-mode-skip-latents — validates SearchUpdater._compute_latent_samples short-circuit in test mode.

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	7	0	1	8
autogalaxy_workspace	6	0	0	6
autolens_workspace	7	0	0	7
autolens_workspace_test	12	0	1	13
euclid_strong_lens_modeling_pipeline	6	0	0	6
HowToLens	6	0	0	6

All 44 smoke tests passed across 6 workspaces (2 pre-existing skips: searches/mcmc in autofit and database/scrape/general in autolens_workspace_test — both already in no_run.yaml).

Notable: cookbooks/latent_variables.py (autofit_workspace) and latent/latent_variables_smoke.py (autolens_workspace_test) both PASS, confirming the latent-variable short-circuit does not break downstream scripts.

[Jammy2211]

Shipped — 2026-05-24

PRs

• PyAutoConf Feature/sample table #109 — feat: add skip_latents() to test_mode helpers PyAutoConf#109 (merged 13:21 UTC)
• PyAutoFit fix: skip _compute_latent_samples in PYAUTO_TEST_MODE (#1294) #1295 — fix: skip _compute_latent_samples in PYAUTO_TEST_MODE (#1294) #1295 (merged 13:21 UTC)

Summary

Added autoconf.test_mode.skip_latents() helper composing is_test_mode() OR PYAUTO_SKIP_LATENTS=1, and short-circuited SearchUpdater._compute_latent_samples to return None when it fires. Eliminates the 100-sample post-fit latent dispatch loop in test-mode runs, which under PyAutoGalaxy's AnalysisDataset.LATENT_BATCH_MODE = "jit" override was ~750s per SLaM fit on CPU. End-to-end SLaM slam_start_here.py smoke went from > 5-min timeout to 29.6s.

Session Notes

• Out of scope, queued as follow-up: the deeper architectural fix — splitting LATENT_BATCH_MODE per-latent so fast latents vmap while only einstein_radius stays jit-per-sample. That's the real-mode SLaM perf fix; this PR only unblocked smoke testing. Spans PyAutoFit dispatcher (_compute_latent_samples) + PyAutoGalaxy registry (LATENT_FUNCTIONS needs to carry a per-entry mode tag) + tests. Filed as a TBA prompt for PyAutoPrompt/autogalaxy/latent_batch_mode_per_latent.md.
• New env var PYAUTO_SKIP_LATENTS=1 is now available alongside PYAUTO_TEST_MODE. Auto-fires inside any test-mode run via skip_latents() = is_test_mode() or PYAUTO_SKIP_LATENTS=1. Discoverable via autoconf/test_mode.py docstring; not added to the fast-smoke env set in workspace CLAUDE.md files because PYAUTO_TEST_MODE already covers smoke.
• Validation: PyAutoConf 95/95, PyAutoFit 1413/1414 (1 skip), new test_autofit/non_linear/search/test_updater.py 3/3, full smoke 44/44 across all 6 workspaces, end-to-end SLaM 29.6s.
• Dependency order mattered: PyAutoConf merged first since PyAutoFit imports the new helper at module load. Both labelled pending-release.