fix: sampler CLI prints variable name, JAX compile msg fires early

[Jammy2211]

Overview

Two things the CLI tells a user during a non-linear search are wrong, and both cost the user the same thing: they cannot tell what the run is doing, or how long a silence will last.

1. A variable name reaches the terminal instead of its value. 23 workspace scripts print On-the-fly updates every iterations_per_quick_update are printed to the notebook. — the literal token, not the cadence.
2. The JAX compile message fires before any compilation happens. autofit/non_linear/fitness.py:508/528/549 log at wrapper-construction time, but jax.jit(...) returns instantly; tracing, lowering and XLA compilation all happen on the first call to the returned function (fitness.py:310). The user sees "applied in 0.0002 seconds", then sits through an unexplained wait that can run to minutes.

A complication found while scoping: the packaged default for iterations_per_quick_update is 1e99 (config/general.yaml, updates: block) — the codebase's documented inf-like "never" sentinel (see the _steps_until_full_update docstring, abstract_search.py:1040-1088). Only hpc_mode sets a real value (250000). So for almost every user the truthful message is not a number at all: it is "these updates are off". A naive f-string would render 1e+99, which is worse than the placeholder it replaces.

Decision: the library owns the cadence message. PyAutoFit logs it once at search start; the 23 workspace scripts drop that sentence entirely — one source of truth instead of the same conditional pasted into 23 files. This issue is Phase 1 (PyAutoFit); the workspace sweep is Phase 2, a separate follow-up issue.

Plan

• Add an AbstractSearch.quick_update_message property that renders the real integer cadence, or — when the value is the 1e99 "never" sentinel — states plainly that on-the-fly updates are disabled and names the config key that enables them.
• Log that message once at search start, alongside the existing "Starting non-linear search…" line.
• Replace the three misleading JAX pre-wrap log pairs with a one-shot wrapper that logs JAX jit compiling <what>, could take seconds or minutes... on the first call — where the wait actually is — and reports the true elapsed compile time after it.
• Cover both with unit tests (neither needs JAX installed).
• File Phase 2 (remove the now-duplicated sentence from 23 workspace scripts, regenerate notebooks/markdown) as a follow-up.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary, Phase 1 — this issue)
• autolens_workspace / autogalaxy_workspace / HowToLens (Phase 2, separate issue)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autolens_workspace	main	clean
./autogalaxy_workspace	main	clean
./HowToLens	main	clean

Suggested branch: feature/sampler-cli-output-numbers

Phase 2 is additionally gated on contention: autolens_workspace currently has two live worktree claims (multi-galaxy-slam-followup, scaling-relation-degraded-profile-fixes).

Brain scored this large (8) and wanted a phase split. The score comes from its repo-count proxy; the library change is one property plus one wrapper. Size overridden to medium, split kept.

Implementation Steps

1a — cadence message. autofit/non_linear/search/abstract_search.py

Add a module constant:

#: Values at or above this disable periodic updates. The packaged config
#: default is the inf-like ``1e99`` sentinel documented on
#: ``_steps_until_full_update`` — kept as a float so ``search.json`` stores a
#: readable ``1e99`` rather than a 99-digit integer.
ITERATIONS_NEVER = 1e90

and a property on AbstractSearch:

@property
def quick_update_message(self) -> str:
    n = self.iterations_per_quick_update
    if not np.isfinite(n) or n >= ITERATIONS_NEVER:
        return (
            "On-the-fly updates of the maximum likelihood model are disabled. "
            "Set `updates: iterations_per_quick_update` in config/general.yaml "
            "to a finite number of iterations to enable them."
        )
    return (
        "On-the-fly updates of the maximum likelihood model every "
        f"{int(n)} iterations."
    )

np is already imported (abstract_search.py:6). Log it in fit() immediately after the existing "Starting non-linear search…" block (abstract_search.py:496-510):

logger.info(self.quick_update_message)

Unconditional, matching the adjacent call — not gated on self.silence.

1b — JAX compile message. autofit/non_linear/fitness.py

Add a one-shot wrapper helper on Fitness:

def _log_on_first_compile(self, func, description):
    """Wrap ``func`` so its first invocation reports the JAX compile it triggers.

    ``jax.jit`` / ``jax.vmap`` / ``jax.grad`` return immediately — tracing,
    lowering and XLA compilation all happen on the first call to the returned
    function. This is therefore the only point at which the user can be told
    why the run appears to hang.
    """
    state = {"compiled": False}

    def wrapper(*args, **kwargs):
        if state["compiled"]:
            return func(*args, **kwargs)
        logger.info(
            f"JAX jit compiling {description}, could take seconds or minutes..."
        )
        start = time.time()
        try:
            result = func(*args, **kwargs)
            # JAX dispatches asynchronously; block once so the reported time is
            # the real wall-clock wait rather than the dispatch latency.
            try:
                import jax
                jax.block_until_ready(result)
            except Exception:
                pass
        finally:
            state["compiled"] = True
        logger.info(
            f"JAX jit compilation of {description} complete in "
            f"{time.time() - start:.1f} seconds."
        )
        return result

    return wrapper

Then collapse the three cached_property bodies (currently fitness.py:494-553) to use it, deleting the misleading pre-wrap timing logs:

• _vmap → jax.vmap(jax.jit(self.call)), description "vectorized (vmap) likelihood function"
• _jit → jax.jit(self.call), description "likelihood function"
• _grad → jax.grad(self.call), description "likelihood function gradient"

Constraints that shaped this (all verified against the source):

• fitness._jit is passed to scipy as fun= in autofit/non_linear/search/mle/bfgs/search.py:183 — the wrapper must stay a plain callable, which it is.
• __getstate__ / __setstate__ (fitness.py:474-491) already strip and lazily rebuild _call/_jit/_vmap/_grad, so the non-picklable closure never reaches pickle and the flag resets per process — correct, since an unpickled process really does recompile. Keeps test_autofit/non_linear/test_fitness_jax_dispatch.py green.
• _warmup_visualization (fitness.py:177-200) already logs "Warming up visualization (one-time JAX compilation)…" — left as is; the new wording is consistent with it.

1c — tests. test_autofit/non_linear/

• _log_on_first_compile: wrap a plain lambda, assert the "could take seconds or minutes" line is logged exactly once across two calls (via caplog) and the return value passes through. No JAX needed.
• quick_update_message: parametrize a finite cadence (renders the integer, no .0) against the 1e99 sentinel (renders the disabled wording naming the config key).

Key Files

• autofit/non_linear/search/abstract_search.py — ITERATIONS_NEVER, quick_update_message, the fit() log call (~line 510)
• autofit/non_linear/fitness.py — _log_on_first_compile, and the _jit / _vmap / _grad cached properties
• autofit/non_linear/search/mle/bfgs/search.py:183 — consumer of fitness._jit, constrains the wrapper to a plain callable
• test_autofit/non_linear/test_fitness_jax_dispatch.py — existing pickle-roundtrip coverage that must stay green

Verification

cd PyAutoFit
NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib \
  python -m pytest test_autofit/

Then end-to-end, which is what actually proves the fix — unit tests cannot see message ordering against a real compile:

• Run a JAX-backed fit and confirm the compile line appears before the long wait, and that the reported elapsed time matches the observed wait rather than ~0s.
• Same run under the default config: confirm the "updates are disabled" wording naming iterations_per_quick_update.
• Re-run with iterations_per_quick_update set to a small finite value: confirm the printed integer matches the config and that quick updates then actually occur at that cadence.

Original Prompt

Click to expand starting prompt

On-the-fly updates every iterations_per_quick_update is the command line output
when a sampler runs, the actuallyh number should be printed to the CLI should say
the actual number

follow up, when JAX compiles a LH function it should say Should say "JAX jit
compiling likelihood function, could take seconds or minutes..." in output so user
knows why theres a small wait so a user knows they are waiting

[Jammy2211]

quick_update_message gives inert advice for the multi-start gradient searches

Surfaced from #1433 (per-step progress logging for MultiStartProdigy / MultiStartAdam / MultiStartADABelief / MultiStartLion), which runs in parallel with this task on disjoint PyAutoFit files. Flagging rather than fixing, since this file is yours.

The new NonLinearSearch.quick_update_message property logs, on the disabled branch:

On-the-fly updates of the maximum likelihood model are disabled. Set updates: iterations_per_quick_update in config/general.yaml to a finite number of iterations to enable them.

That is correct and actionable for the samplers. For the four multi-start gradient searches it is inert advice — a user who follows it changes the config, sees no updates, and has no way to learn why. Two independent reasons, either one sufficient:

1. The searches never wire the knob into their Fitness.

AbstractMultiStartGradient._fit constructs it without the kwarg (autofit/non_linear/search/mle/multi_start_gradient/search.py:234-241):

fitness = Fitness(
    model=model,
    analysis=analysis,
    paths=self.paths,
    fom_is_log_likelihood=False,
    resample_figure_of_merit=-np.inf,
    convert_to_chi_squared=True,
)

Every other search passes it through — e.g. nest/nautilus/search.py:195,216 use iterations_per_quick_update=self.iterations_per_quick_update. Since Fitness.__init__ defaults it to None (fitness.py:47), Fitness.quick_update() hits its if self.iterations_per_quick_update is None: return guard (fitness.py:378) on every call, whatever config/general.yaml says.

2. Even wired, the counter cannot run.

The counting lives inside fitness.call, which these searches trace under jax.jit(jax.vmap(...)) (search.py:260-261). The Python-side self.quick_update_count += total_updates executes once at trace time and never again — and if best_log_likelihood > self.quick_update_max_lh: on a traced value would not survive as a Python branch anyway.

Suggested handling

The honest message for these searches is "not available", not "disabled — enable it here". Options, cheapest first:

• Have quick_update_message consult whether the search actually wires the knob through (e.g. a class-level flag such as supports_quick_update = False on AbstractMultiStartGradient), and emit a third branch naming the real reason — the objective is JIT-traced, so per-iteration Python callbacks cannot fire.
• Or leave the wording alone here and let feat: per-step progress logging for the multi-start gradient searches #1433 override the property for the gradient searches.

Either way it belongs on this issue's side of the seam, since #1433 does not touch abstract_search.py.

For context: #1433 is adding per-step progress logging plus two JAX-compile notices directly inside _fit, because these searches bypass fitness._jit / _vmap / _grad entirely — so the compile-message fix planned here will not reach them either. Detail in #1433 (comment). Wording will be matched to whatever this task lands.

[Jammy2211]

Scope grew during verification — two real defects found behind the cosmetic one

The planned output fixes are done, but verifying them end-to-end (rather than only by unit test) turned up two substantive bugs that the cosmetic fix would otherwise have made worse. Both are now fixed here, with the wiring expansion human-approved before it was made.

1. Quick updates were silently dead for 7 of 8 searches

Only Nautilus forwarded iterations_per_quick_update when constructing its Fitness. Dynesty, Emcee, Zeus, BlackJAX NUTS, BFGS, Drawer and MultiStartGradient all omitted it, so Fitness.iterations_per_quick_update was None and manage_quick_update returned at its first guard (fitness.py:378). The search ran perfectly and simply never updated.

Control test — same model, same cadence of 25, fresh output directory, only the search differs:

Search	forwards the kwarg	startup message	updates fire
Nautilus	yes	"every 25 iterations"	✅
DynestyStatic	no	"every 25 iterations"	❌ never

So announcing the cadence in the startup log would have turned a dead feature into an active lie. Fixed by forwarding iterations_per_quick_update / background_quick_update / live_visual_update at the 6 remaining search construction sites. Before/after on identical input: 0 quick updates → 1.

multi_start_gradient/search.py is deliberately not wired — it differentiates fitness.call inside its own jit/vmap step loop, so the Python-side counter would run once at trace time. Its progress reporting is #1433's. The exemption is recorded with that reason in the new wiring test.

2. Wiring it up exposed a latent crash in manage_quick_update

File "autofit/non_linear/fitness.py", line 429, in manage_quick_update
    max_log_likelihood_sample=self.quick_update_max_lh_parameters.tolist(),
AttributeError: 'list' object has no attribute 'tolist'

Searches hand over parameters in whatever type they hold: ndarray (Nautilus), JAX array, or a plain Python list (Dynesty's initializer). .tolist() assumed the array case — which held only while Nautilus was the sole caller. Normalized with np.asarray(...).tolist().

3. Same misleading-message bug in the latent path

analysis/latent.py had the identical pattern — "may take a few seconds", reported 0.0011 seconds, then compiled later. The helper is now a shared autofit/non_linear/jax_compile.py (log_on_first_compile) used by both fitness.py and latent.py, rather than duplicated.

Verified

Real DynestyStatic + JAX fit, clean output directory:

Starting non-linear search with JAX (CPU: cpu).
On-the-fly updates of the maximum likelihood model every 25 iterations.
JAX jit compiling likelihood function, could take seconds or minutes...
JAX jit compilation of likelihood function complete in 0.2 seconds.
Performing quick update of maximum log likelihood fit image and model.results   (x3)
...
JAX jit compiling latent variable function (vmap), could take seconds or minutes...

Isolated probe against a genuinely expensive jax.jit function: nothing logged at wrap time; reported 0.7s matched the real 0.7s wall-clock; second call silent and 64x faster.

Both new detectors were checked against the pre-fix source and do fail there — the wiring test on 6 sites, the parameter-container test on the list and tuple cases.

Full suite: 1615 passed, 2 skipped.

Default config still reports On-the-fly updates ... are disabled. Set 'updates: iterations_per_quick_update' in config/general.yaml ..., since the packaged default is the 1e99 never-sentinel.

[Jammy2211]

Library PR Created

PyAutoFit#1436 — #1436 (OPEN, pending-release)

Heart gated YELLOW, acknowledged by the human before shipping. The three reasons were unrelated to this change: manifest drift: tenant firewall (organ code) — 1 mismatch(es) vs PyAutoMind/repos.yaml, plus two stale markers (test run status unknown (no report.json), release validation stale: source moved since rehearsal). No RED reasons. PyAutoFit's own suite is green at 1615 passed / 2 skipped.

Workspace impact

Option (i)/(ii)-adjacent: no API migration is required, but 23 workspace scripts now carry a sentence the library emits itself, and theirs still names the variable instead of the value. Phase 2 removes it.

Repo	scripts
autolens_workspace	14+
autogalaxy_workspace	3+
HowToLens	1

Plus 23 generated .ipynb and 14 .md to regenerate.

Phase 2 prompt is filed at PyAutoMind/draft/docs/autolens_workspace/sampler_cli_output_workspace_sweep.md and should not start until #1436 merges — until then the workspace sentence is wrong but harmless, whereas removing it early would leave users with no cadence message at all.

Mind state: library-shipped, workspace-pending.

[Jammy2211]

CI green

PR #1436 — all three legs pass on daa0dbc:

Check	Result
unittest / unittest (3.12)	pass (3m01s)
unittest / unittest (3.13)	pass (3m06s)
docs / docs-build	pass (1m12s)

PR state MERGEABLE / CLEAN. Awaiting human merge — Phase 2 (the 23 workspace scripts) stays parked until it lands.

[Jammy2211]

Shipped

PyAutoFit#1436 MERGED 2026-07-30 (merge commit 39cbce779). CI green on all three legs before merge; issue auto-closed.

Cleanup done: worktree removed, feature/sampler-cli-output-numbers deleted locally and on origin, canonical PyAutoFit synced to main.

Completion record: PyAutoMind/complete/2026/07/sampler-cli-output-numbers.md (lifecycle check: OK, index refreshed).

Follow-up still open

Phase 2 — remove the now-duplicated sentence from 23 workspace scripts (autolens_workspace 14+, autogalaxy_workspace 3+, HowToLens 1) and regenerate the 23 .ipynb / 14 .md counterparts. Prompt filed at PyAutoMind/draft/docs/autolens_workspace/sampler_cli_output_workspace_sweep.md; needs its own issue via start_dev. Now unblocked, since the library side has merged.

Note for #1433: the compile-message wording that landed is JAX jit compiling <what>, could take seconds or minutes... via the shared autofit/non_linear/jax_compile.py helper (log_on_first_compile). Reuse it rather than adding a second phrasing.