feat: re-mainline af.NSS as a first-class search on mainline blackjax ≥1.6

autofit/__init__.py

@@ -95,6 +95,7 @@
 from .non_linear.search.mcmc.emcee.search import Emcee
 from .non_linear.search.mcmc.zeus.search import Zeus
 from .non_linear.search.nest.nautilus.search import Nautilus
+from .non_linear.search.nest.nss.search import NSS
 from .non_linear.search.nest.dynesty.search.dynamic import DynestyDynamic
 from .non_linear.search.nest.dynesty.search.static import DynestyStatic
 from .non_linear.search.mle.drawer.search import Drawer

autofit/non_linear/search/nest/nss/__init__.py

@@ -0,0 +1 @@
+from .search import NSS

autofit/non_linear/search/nest/nss/_chunked_nss.py

@@ -0,0 +1,138 @@
+"""Local replica of ``blackjax.ns.nss.as_top_level_api`` with chunked vmap sites.
+
+PyAutoFit#1303 added ``make_chunked_update_strategy`` for the per-iteration
+MCMC step path inside ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``.
+That covers the inner vmap over ``num_delete`` particles, but **not** the
+separate hardcoded ``jax.vmap(init_state_fn)`` inside
+``blackjax.ns.nss.as_top_level_api``'s ``init_fn`` — and that's where
+inversion-heavy lensing cells (PyAutoLens pixelization / Delaunay at HST
+scale) OOM on A100 80 GB before the sampling loop even starts.
+
+This module replaces ``blackjax.nss(...)`` with a local builder so we
+control both seams:
+
+- The chunked update strategy replaces the hardcoded
+  ``update_with_mcmc_take_last`` inside the kernel composition (mainline
+  blackjax >= 1.6 exposes no ``update_strategy=`` kwarg — the pre-1.6 fork
+  did — so the kernel is composed here from mainline's public helpers:
+  ``build_slice_kernel`` / ``slice_constrained_step`` /
+  ``build_adaptive_kernel``).
+- The chunked init swaps ``jax.vmap(init_state_fn)`` for
+  ``jax.lax.map(init_state_fn, positions, batch_size=chunk_size)`` so peak
+  GPU memory during ``algo.init`` becomes ``chunk_size × per_particle_state``
+  instead of ``n_live × per_particle_state``.
+
+When ``chunk_size`` is None the builder still uses ``jax.vmap`` and is
+bit-identical to upstream. The builder otherwise produces a
+``blackjax.SamplingAlgorithm`` with the same shape ``blackjax.nss(...)``
+returns, so ``af.NSS._fit`` is a one-line switch.
+
+See PyAutoFit#1304 for the diagnosis and A100 evidence (jobs 322605 /
+322606 OOM at the same byte counts as before #1303 landed, because the
+crash is in ``algo.init`` not ``algo.step``).
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional
+
+
+def build_chunked_nss_algorithm(
+    *,
+    logprior_fn: Callable,
+    loglikelihood_fn: Callable,
+    num_inner_steps: int,
+    num_delete: int,
+    chunk_size: Optional[int],
+):
+    """Return a ``blackjax.SamplingAlgorithm`` with chunked init + step paths.
+
+    Replicates the body of ``blackjax.ns.nss.as_top_level_api`` (the
+    handley-lab fork) so we can plug in chunked variants of the two
+    vmap sites — the inner MCMC step (via the existing
+    ``make_chunked_update_strategy``) and the n_live-wide init.
+
+    Parameters
+    ----------
+    logprior_fn
+        Log-prior callable, ``positions -> scalar log-prior``.
+    loglikelihood_fn
+        Log-likelihood callable, ``positions -> scalar log-L``.
+    num_inner_steps
+        Number of HRSS steps per particle replacement (matches
+        ``af.NSS.num_mcmc_steps``).
+    num_delete
+        Number of particles replaced per outer iteration (matches
+        ``af.NSS.num_delete``).
+    chunk_size
+        Optional GPU-memory knob. When None, both vmap sites use plain
+        ``jax.vmap`` and the result is bit-identical to upstream
+        ``blackjax.nss(...)``. When set, peak memory in each site becomes
+        ``chunk_size × per_particle_state``.
+    """
+    # Local imports keep this module cheap to import when ``af.NSS`` is
+    # never used (blackjax + jax are optional deps).
+    import jax
+    from blackjax import SamplingAlgorithm
+    from blackjax.ns.adaptive import init as ns_init
+    from blackjax.ns.base import init_state_strategy
+    from blackjax.ns.from_mcmc import build_adaptive_kernel, default_delete_fn
+    from blackjax.ns.nss import (
+        build_slice_kernel,
+        covariance_proposal,
+        live_covariance,
+        slice_constrained_step,
+        stepping_out,
+    )
+
+    from ._chunked_update import (
+        make_chunked_update_strategy,
+    )
+
+    init_state_fn = partial(
+        init_state_strategy,
+        logprior_fn=logprior_fn,
+        loglikelihood_fn=loglikelihood_fn,
+    )
+
+    # Replicates mainline ``blackjax.ns.nss.build_kernel`` +
+    # ``ns.from_mcmc.build_kernel`` with the hardcoded
+    # ``update_with_mcmc_take_last`` swapped for the chunked variant.
+    # (The pre-1.6 fork exposed an ``update_strategy=`` kwarg for this;
+    # mainline does not, hence the explicit composition.)
+    slice_kernel = build_slice_kernel(
+        interval=stepping_out,
+        max_expansions=10,
+        max_shrinkage=100,
+    )
+    constrained_step_fn = slice_constrained_step(
+        init_state_fn, slice_kernel, covariance_proposal
+    )
+    inner_kernel = make_chunked_update_strategy(chunk_size)(
+        constrained_step_fn, num_inner_steps, num_delete
+    )
+    kernel = build_adaptive_kernel(
+        partial(default_delete_fn, num_delete=num_delete),
+        inner_kernel,
+        update_inner_kernel_params_fn=live_covariance,
+    )
+
+    def init_fn(position, rng_key=None):
+        if chunk_size is None:
+            init_batcher = jax.vmap(init_state_fn)
+        else:
+            init_batcher = lambda p: jax.lax.map(
+                init_state_fn, p, batch_size=chunk_size
+            )
+        return ns_init(
+            position,
+            init_state_fn=init_batcher,
+            update_inner_kernel_params_fn=live_covariance,
+            rng_key=rng_key,
+        )
+
+    def step_fn(rng_key, state):
+        return kernel(rng_key, state)
+
+    return SamplingAlgorithm(init_fn, step_fn)

autofit/non_linear/search/nest/nss/_chunked_update.py

@@ -0,0 +1,117 @@
+"""Chunked replacement for ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``.
+
+Upstream blackjax fans out ``num_delete`` particles through ``jax.vmap``
+with no chunking:
+
+    sample_keys = jax.random.split(sample_key, num_delete)
+    return jax.vmap(mcmc_kernel)(sample_keys, start_state)
+
+On inversion-heavy likelihoods (e.g. PyAutoLens pixelization / Delaunay
+source models) the per-particle MCMC state plus scatter temp buffers
+exceeds A100 80 GB even at ``num_delete=16``. See PyAutoFit#1301 for the
+full diagnosis and per-cell evidence from ``autolens_profiling``.
+
+``chunked_update_with_mcmc_take_last`` accepts a ``chunk_size`` kwarg and
+swaps the vmap for ``jax.lax.map(..., batch_size=chunk_size)`` when
+``chunk_size < num_delete`` — same vmap parallelism within a chunk, sequential
+chunks across. Peak memory becomes ``chunk_size × per_particle_state``
+instead of ``num_delete × per_particle_state``.
+
+When ``chunk_size`` is None or ``>= num_delete`` the function is
+bit-identical to upstream.
+
+Mainline blackjax (>= 1.6) hardcodes ``update_with_mcmc_take_last``
+inside ``blackjax.ns.from_mcmc.build_kernel`` — there is no
+``update_strategy=`` kwarg (the pre-1.6 fork had one). The chunked
+variant is therefore wired in by ``_chunked_nss.build_chunked_nss_algorithm``,
+which recomposes the kernel from mainline's public helpers.
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional
+
+
+def make_chunked_update_strategy(chunk_size: Optional[int]) -> Callable:
+    """Return a chunked replacement for ``update_with_mcmc_take_last``.
+
+    Signature matches ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``
+    so ``_chunked_nss.build_chunked_nss_algorithm`` can drop it into the
+    kernel composition unmodified.
+
+    Parameters
+    ----------
+    chunk_size
+        Number of particles to vmap-batch per chunk. When None or
+        ``>= num_delete`` the chunked path is skipped and the function
+        falls through to a plain ``jax.vmap`` (matching upstream
+        behaviour bit-for-bit).
+    """
+
+    def chunked_update_with_mcmc_take_last(
+        constrained_mcmc_step_fn,
+        num_mcmc_steps,
+        num_delete,
+    ):
+        """Drop-in for ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``.
+
+        Identical to upstream except the inner
+        ``jax.vmap(mcmc_kernel)(sample_keys, start_state)`` is replaced
+        with ``jax.lax.map(..., batch_size=chunk_size)`` when
+        ``chunk_size`` is set and smaller than ``num_delete``.
+        """
+        import jax
+        import jax.numpy as jnp
+
+        def update_function(rng_key, state, loglikelihood_0, **step_parameters):
+            choice_key, sample_key = jax.random.split(rng_key)
+            particles = state.particles
+
+            # Select start particles from survivors (verbatim from upstream).
+            weights = (particles.loglikelihood > loglikelihood_0).astype(jnp.float32)
+            weights = jnp.where(weights.sum() > 0.0, weights, jnp.ones_like(weights))
+            start_idx = jax.random.choice(
+                choice_key,
+                len(weights),
+                shape=(num_delete,),
+                p=weights / weights.sum(),
+                replace=True,
+            )
+            start_state = jax.tree.map(lambda x: x[start_idx], particles)
+
+            shared_mcmc_step_fn = partial(
+                constrained_mcmc_step_fn,
+                loglikelihood_0=loglikelihood_0,
+                **step_parameters,
+            )
+
+            def mcmc_kernel(rng_key, state):
+                keys = jax.random.split(rng_key, num_mcmc_steps)
+
+                def body_fn(state, rng_key):
+                    new_state, info = shared_mcmc_step_fn(rng_key, state)
+                    return new_state, info
+
+                final_state, infos = jax.lax.scan(body_fn, state, keys)
+                return final_state, infos
+
+            sample_keys = jax.random.split(sample_key, num_delete)
+
+            # Fall through to bit-identical upstream behaviour when the
+            # user hasn't asked for chunking, or when the requested chunk
+            # already covers every particle.
+            if chunk_size is None or chunk_size >= num_delete:
+                return jax.vmap(mcmc_kernel)(sample_keys, start_state)
+
+            # Chunked path: jax.lax.map(batch_size=k) vmaps within each
+            # chunk-of-k particles and loops across chunks.
+            return jax.lax.map(
+                lambda xs: mcmc_kernel(xs[0], xs[1]),
+                (sample_keys, start_state),
+                batch_size=chunk_size,
+            )
+
+        return update_function
+
+    return chunked_update_with_mcmc_take_last

autofit/non_linear/search/nest/nss/samples.py

@@ -0,0 +1,25 @@
+from autofit.non_linear.samples.nest import SamplesNest
+
+
+class NSSamples(SamplesNest):
+    """Posterior samples from an ``af.NSS`` (Nested Slice Sampling) fit.
+
+    Thin subclass of ``SamplesNest`` that exists primarily for type identity —
+    aggregator code and downstream consumers can distinguish an NSS run from
+    a Nautilus / Dynesty run by ``isinstance(samples, NSSamples)``. The
+    actual posterior + log-evidence wiring is inherited from ``SamplesNest``
+    (and ultimately ``SamplesPDF``); ``af.NSS`` builds the ``sample_list`` in
+    ``NSS.samples_via_internal_from``.
+    """
+
+    @property
+    def log_evidence_error(self) -> float:
+        """The stochastic batch-error of the NS evidence estimate.
+
+        NSS returns an array of log-evidence estimates ``logZs`` across the
+        live ensemble (the Monte Carlo simulation of ``log_dX``). This
+        property exposes the standard deviation as the natural per-sample
+        uncertainty. The mean is reported as ``log_evidence`` in
+        ``samples_info``.
+        """
+        return float(self.samples_info.get("log_evidence_error", float("nan")))