fix: skip _compute_latent_samples in PYAUTO_TEST_MODE (#1294)

[Jammy2211]

Summary

Short-circuit SearchUpdater._compute_latent_samples to return None when autoconf.test_mode.skip_latents() is True. Eliminates the post-fit latent-processing cost in test-mode runs.

Root cause. Under PYAUTO_TEST_MODE=2 the sampler is already bypassed (mocks the fit), but the post-fit pass still draws output.latent_draw_via_pdf_size = 100 samples and dispatches analysis.compute_latent_samples for each. With PyAutoGalaxy's AnalysisDataset.LATENT_BATCH_MODE = "jit" override (forced because the einstein-radius latent's jax_zero_contour.ZeroSolver uses lax.cond / lax.while_loop early-termination and is not vmap-compatible), that's a Python loop of 100 JIT dispatches per stage — ~750 s per SLaM fit on CPU. The smoke caller hit the 5-min timeout still inside the first SLaM stage's latent loop.

Fix. Two lines at the top of _compute_latent_samples: import skip_latents from autoconf.test_mode and return None when it fires. Test-mode runs skip the path entirely; real-mode runs are unchanged. PYAUTO_SKIP_LATENTS=1 provides the same bypass for real-mode users who don't want the post-fit latent work.

Validation. SLaM slam_start_here.py end-to-end with the full smoke env (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) completes in 29.6 s — all 5 stages (source_lp[1] → source_pix[1] → source_pix[2] → light[1] → mass_total[1]) run to completion. Previously timed out at 5 min inside source_lp[1] latent processing.

Out of scope. The deeper architectural fix — splitting LATENT_BATCH_MODE per-latent so fast latents vmap while einstein_radius stays jit-per-sample — is a separate follow-up (real-mode SLaM perf). This PR only unblocks smoke testing.

Dependency. Depends on the PyAutoConf PR that adds skip_latents(). Merge that one first.

Fixes #1294.

API Changes

Public API: no change. Purely internal — adds if skip_latents(): return None at the top of the private SearchUpdater._compute_latent_samples method. The PYAUTO_SKIP_LATENTS env var is a new user-facing env knob but documented inline in autoconf.test_mode.skip_latents()'s docstring.

See full details below.

Test Plan

• python -m pytest test_autofit/ — 1413 passed, 1 skipped.
• New test_autofit/non_linear/search/test_updater.py — 3/3 (test-mode trigger, PYAUTO_SKIP_LATENTS=1 trigger, config-gate regression).
• End-to-end SLaM slam_start_here.py with full smoke env completes in 29.6s (previously > 5min, hit timeout).

Full API Changes (for automation & release notes)

Removed / Renamed

None.

Added

• PYAUTO_SKIP_LATENTS=1 environment variable — when set, SearchUpdater._compute_latent_samples short-circuits and returns None. New env-var knob; previously latent post-processing was only controlled by the output.latent_after_fit / output.latent_during_fit config gate.

Changed Behaviour

• SearchUpdater._compute_latent_samples (private) now also short-circuits when autoconf.test_mode.skip_latents() returns True — i.e., any PYAUTO_TEST_MODE level set, or PYAUTO_SKIP_LATENTS=1. Real-mode runs with neither env var set are unchanged.

Migration

None for end users. Internal callers do not need to update — the method's Optional[Samples] return type already accommodated None.

🤖 Generated with Claude Code