feat: warn before an unbatched MultiStartProdigy jvp exhausts memory

autofit/non_linear/analysis/analysis.py

@@ -334,23 +334,35 @@ def supports_jax_visualization(self) -> bool:
     def perform_quick_update(self, paths, instance):
         raise NotImplementedError
 
-    def print_vram_use(self, model, batch_size : int) -> str:
+    def batched_memory_bytes(self, model, batch_size: int, gradient: bool = False):
         """
-        Print JAX VRAM use for a given batch size.
+        Measure the memory a batched evaluation of this analysis would need.
+
+        Compiles (but never runs) ``jax.vmap`` over ``batch_size`` copies of the
+        model at its prior medians and reads XLA's own memory analysis, so this
+        is a lowering-time estimate — it cannot itself OOM the way the run it is
+        predicting would.
 
         Parameters
         ----------
         batch_size
-            The batch size to profile, which is the number of model evaluations JAX will perform simultaneously.
-        """
-        from autofit.non_linear.test_mode import skip_fit_output
-
-        if skip_fit_output():
-            return
+            Number of model evaluations JAX would perform simultaneously.
+        gradient
+            Measure ``jax.value_and_grad`` of the likelihood rather than the
+            likelihood alone. **A gradient search must pass True.** The jvp
+            carries the whole forward tape, so a likelihood-only measurement
+            can under-report a gradient optimizer's footprint by a large
+            factor — that gap is exactly what let a 48-start interferometer fit
+            reach ~86 GB while the likelihood-only figure looked modest
+            (PyAutoFit#1452).
 
+        Returns
+        -------
+        The projected bytes, or ``None`` when this analysis is not on the JAX
+        path (nothing to measure).
+        """
         if not self._use_jax:
-            print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.")
-            return
+            return None
 
         import jax
         import jax.numpy as jnp
@@ -372,22 +384,54 @@ def print_vram_use(self, model, batch_size : int) -> str:
 
         parameters = jnp.array(parameters)
 
-        batched_call = jax.jit(jax.vmap(fitness.call))
+        call = jax.value_and_grad(fitness.call) if gradient else fitness.call
+
+        batched_call = jax.jit(jax.vmap(call))
         lowered = batched_call.lower(parameters)
         compiled = lowered.compile()
         memory_analysis = compiled.memory_analysis()
 
-        vram_bytes = (
-                memory_analysis.output_size_in_bytes
-                + memory_analysis.temp_size_in_bytes
+        return (
+            memory_analysis.output_size_in_bytes
+            + memory_analysis.temp_size_in_bytes
+        )
+
+    def print_vram_use(self, model, batch_size : int, gradient: bool = False) -> str:
+        """
+        Print JAX VRAM use for a given batch size.
+
+        Parameters
+        ----------
+        batch_size
+            The batch size to profile, which is the number of model evaluations JAX will perform simultaneously.
+        gradient
+            Profile ``value_and_grad`` rather than the likelihood alone. Pass
+            True when the search is a gradient optimizer (e.g.
+            ``af.MultiStartProdigy``) — see ``batched_memory_bytes``.
+        """
+        from autofit.non_linear.test_mode import skip_fit_output
+
+        if skip_fit_output():
+            return
+
+        if not self._use_jax:
+            print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.")
+            return
+
+        vram_bytes = self.batched_memory_bytes(
+            model=model, batch_size=batch_size, gradient=gradient
         )
 
+        # Stated so the number is never read as covering more than it measures:
+        # the likelihood-only figure does not bound a gradient search.
+        measured = "likelihood + gradient" if gradient else "likelihood only"
+
         if vram_bytes == 0:
             print(
                 "VRAM USE = 0.000 GB "
                 "(this likely means JAX is running in CPU-only mode)"
             )
         else:
             print(
-                f"VRAM USE = {vram_bytes / 1024 ** 3:.3f} GB"
+                f"VRAM USE = {vram_bytes / 1024 ** 3:.3f} GB ({measured}, batch_size={batch_size})"
             )

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -353,6 +353,96 @@ def _progress_message(
 
         return " | ".join(parts)
 
+    def _memory_budget_bytes(self):
+        """
+        Bytes available to the batched jvp, or ``None`` if not determinable.
+
+        Prefers the JAX device's own limit (the GPU case this guard mainly
+        exists for) and falls back to free host RAM.
+        """
+        try:
+            import jax
+
+            stats = jax.devices()[0].memory_stats() or {}
+            for key in ("bytes_limit", "bytes_reservable_limit"):
+                if stats.get(key):
+                    return int(stats[key])
+        except Exception:
+            pass
+
+        try:
+            import psutil
+
+            return int(psutil.virtual_memory().available)
+        except Exception:
+            return None
+
+    def _warn_if_unbatched_exceeds_memory(self, model, analysis):
+        """
+        Warn, before compiling the full ``n_starts`` vmap, when its batched jvp
+        is projected to exceed available memory — naming the ``batch_size`` that
+        would fit.
+
+        Without this the only signal is XLA's ``RESOURCE_EXHAUSTED: Out of
+        memory allocating N bytes``, raised from inside the compiled program
+        with nothing pointing at the knob that fixes it. That cost two nightly
+        release runs in 2026-07 (PyAutoFit#1452).
+
+        This only ever *warns*. ``batch_size`` is numerically inert, so
+        auto-applying the suggestion would be safe in principle, but silently
+        changing the execution shape of every existing run on the strength of a
+        projection is not a trade this should make on its own.
+
+        **Known limitation.** ``memory_analysis`` reports 0 on a CPU-only JAX
+        build, which is where the release harness runs. The projection is
+        skipped entirely in that case rather than guessing, so this guard
+        currently helps GPU users and leaves the CPU path to the (now
+        unfiltered) traceback. Making the CPU path measurable is follow-up
+        work and needs validating against a real run, not a unit test.
+
+        Any failure here is swallowed: a memory projection must never be the
+        reason a fit does not start.
+        """
+        try:
+            probe = getattr(analysis, "batched_memory_bytes", None)
+            if probe is None or self.n_starts is None or self.n_starts <= 1:
+                return
+
+            bytes_at_1 = probe(model=model, batch_size=1, gradient=True)
+            if not bytes_at_1:  # None (not on JAX) or 0 (CPU-only: unmeasurable)
+                return
+
+            bytes_at_2 = probe(model=model, batch_size=2, gradient=True)
+            if not bytes_at_2:
+                return
+
+            budget = self._memory_budget_bytes()
+            if not budget:
+                return
+
+            fixed, per_start = batch_memory_model(bytes_at_1, bytes_at_2)
+            suggested = batch_size_within_budget(
+                fixed, per_start, self.n_starts, budget
+            )
+            if suggested is None:
+                return
+
+            projected = fixed + self.n_starts * per_start
+            gb = 1024 ** 3
+            self.logger.warning(
+                f"{type(self).__name__} is set to evaluate all "
+                f"{self.n_starts} starts in one jax.vmap (batch_size=None). "
+                f"Its batched value_and_grad is projected to need "
+                f"~{projected / gb:.1f} GB against ~{budget / gb:.1f} GB "
+                f"available, so this fit is likely to fail with an XLA "
+                f"RESOURCE_EXHAUSTED error. Pass batch_size={suggested} (or "
+                f"lower) to sweep the starts in chunks instead — the result is "
+                f"numerically identical, only the allocation and dispatch "
+                f"change."
+            )
+        except Exception:
+            return
+
     def _fit(
         self,
         model: AbstractPriorModel,
@@ -416,6 +506,7 @@ def _fit(
         _vmapped = jax.jit(jax.vmap(_value_and_grad))
 
         if self.batch_size is None:
+            self._warn_if_unbatched_exceeds_memory(model=model, analysis=analysis)
             batched_value_and_grad = _vmapped
         else:
             batch_size = self.batch_size
@@ -879,6 +970,61 @@ def samples_via_internal_from(
         )
 
 
+def batch_memory_model(bytes_at_1, bytes_at_2):
+    """
+    Split a measured batched-jvp footprint into its fixed and per-start parts.
+
+    Two probe points are enough because the footprint is affine in the batch
+    width: ``bytes(n) = fixed + n * per_start``. ``per_start`` is the slope
+    (what one extra start costs), ``fixed`` the intercept (buffers shared by
+    the whole batch, e.g. the dataset and any persistent operator).
+
+    Measuring *both* is the point. Guidance in the workspaces has claimed VRAM
+    "does not scale with batch size for the persistent buffers, so if it fits
+    at ``batch_size=1`` you can push it up" — true only when ``per_start`` is
+    small next to ``fixed``. For an interferometer likelihood under
+    ``value_and_grad`` it is not: the 2026-07-30 release failure projected to
+    ~1.79 GB *per start*, so 48 starts wanted ~86 GB and the single-start
+    probe that "fit" said nothing useful (PyAutoFit#1452).
+
+    Returns ``(fixed, per_start)``, both clamped at 0 — a non-monotonic or
+    noisy pair of measurements must never yield a negative slope and so a
+    nonsensically large batch suggestion.
+    """
+    per_start = max(0, bytes_at_2 - bytes_at_1)
+    fixed = max(0, bytes_at_1 - per_start)
+    return fixed, per_start
+
+
+def batch_size_within_budget(fixed, per_start, n_starts, budget_bytes):
+    """
+    The largest batch width whose projected footprint fits ``budget_bytes``.
+
+    Returns ``None`` when the full ``n_starts`` batch already fits — the
+    caller should then leave ``batch_size=None`` and keep today's single-vmap
+    fast path. Otherwise returns an int in ``[1, n_starts)``.
+
+    A budget that cannot fit even one start still returns 1: tiling to a
+    single start is the most this knob can do, and a hard failure is the
+    search's to report, not this projection's to pre-empt.
+    """
+    if budget_bytes <= 0 or n_starts <= 1:
+        return None
+
+    def projected(n):
+        return fixed + n * per_start
+
+    if projected(n_starts) <= budget_bytes:
+        return None
+
+    if per_start <= 0:
+        # Fixed cost alone blows the budget; tiling cannot help, but a
+        # single start is still the smallest thing we can ask for.
+        return 1
+
+    return max(1, min(n_starts - 1, int((budget_bytes - fixed) // per_start)))
+
+
 def _chunk_slices(n_rows, batch_size):
     """
     The ``(lo, hi, pad)`` chunk bounds the batched ``value_and_grad`` sweep

test_autofit/non_linear/search/mle/test_multi_start_gradient.py

@@ -6,7 +6,11 @@
 import autofit as af
 from autofit import example
 from autofit.non_linear.search import abstract_search
-from autofit.non_linear.search.mle.multi_start_gradient.search import _chunk_slices
+from autofit.non_linear.search.mle.multi_start_gradient.search import (
+    _chunk_slices,
+    batch_memory_model,
+    batch_size_within_budget,
+)
 from autonerves.dictable import from_dict, to_dict
 
 # The MultiStart gradient searches are JAX-native at fit time, but their
@@ -691,3 +695,71 @@ def test__samples_info__reports_a_cleared_stop_reason_as_unfinished():
 
     assert samples.samples_info["stop_reason"] is None
     assert samples.samples_info["converged"] is False
+
+
+# --- Batched-jvp memory projection (PyAutoFit#1452) ---------------------------
+#
+# Same contract as the _chunk_slices tests above: the projection arithmetic is
+# pure Python and lives here, while the JAX measurement it consumes
+# (`Analysis.batched_memory_bytes`) is exercised in autofit_workspace_test.
+
+GB = 1024 ** 3
+
+# The 2026-07-30 release failure: XLA reported 85,898,814,480 bytes for a
+# 48-start interferometer jvp, i.e. 1,789,558,635 bytes per start.
+INCIDENT_PER_START = 1_789_558_635
+
+
+def test__batch_memory_model__recovers_fixed_and_per_start():
+    fixed = 2 * GB
+    at_1 = fixed + INCIDENT_PER_START
+    at_2 = fixed + 2 * INCIDENT_PER_START
+
+    assert batch_memory_model(at_1, at_2) == (fixed, INCIDENT_PER_START)
+
+
+@pytest.mark.parametrize(
+    "at_1, at_2",
+    [
+        (10, 4),  # non-monotonic measurements
+        (10, 10),  # identical measurements
+    ],
+)
+def test__batch_memory_model__never_yields_a_negative_slope(at_1, at_2):
+    # A negative slope would invert the budget arithmetic and suggest an
+    # absurdly large batch, which is worse than saying nothing.
+    fixed, per_start = batch_memory_model(at_1, at_2)
+    assert per_start == 0
+    assert fixed >= 0
+
+
+def test__batch_size_within_budget__suggests_the_largest_batch_that_fits():
+    fixed, per_start = 2 * GB, INCIDENT_PER_START
+
+    suggested = batch_size_within_budget(fixed, per_start, 48, 16 * GB)
+
+    assert 1 <= suggested < 48
+    assert fixed + suggested * per_start <= 16 * GB
+    assert fixed + (suggested + 1) * per_start > 16 * GB
+
+
+def test__batch_size_within_budget__none_when_the_full_batch_fits():
+    # None means "leave batch_size=None" — the single-vmap fast path is kept.
+    assert batch_size_within_budget(2 * GB, INCIDENT_PER_START, 48, 400 * GB) is None
+
+
+@pytest.mark.parametrize(
+    "fixed, per_start, n_starts, budget, expected",
+    [
+        (100 * GB, 0, 48, 1 * GB, 1),  # fixed cost alone busts it; tiling cannot help
+        (1, 0, 48, 1 * GB, None),  # no per-start cost and it fits
+        (0, 1, 48, 47, 47),  # never returns n_starts itself
+        (0, 100 * GB, 48, 1 * GB, 1),  # not even one start fits -> still 1
+        (2 * GB, INCIDENT_PER_START, 48, 0, None),  # unknown budget
+        (2 * GB, INCIDENT_PER_START, 1, 1, None),  # nothing to tile
+    ],
+)
+def test__batch_size_within_budget__degenerate_inputs(
+    fixed, per_start, n_starts, budget, expected
+):
+    assert batch_size_within_budget(fixed, per_start, n_starts, budget) == expected