Inference API

The pyrox.inference subpackage exposes ensemble-of-MAP and ensemble-of-VI runners as a layered surface — pick the level of control that fits your use case.

Layer 1 — Functional primitives

Roll your own training loop on top of the vmapped state primitives.

ensemble_init(init_fn: Callable[[PRNGKeyArray], PyTree], optimizer: optax.GradientTransformation, *, ensemble_size: int, seed: PRNGKeyArray) -> EnsembleState

Initialize an ensemble of (params, opt_state) by vmap over keys.

Parameters:

Name	Type	Description	Default
init_fn	Callable[[PRNGKeyArray], PyTree]	key -> params. Called once per ensemble member.	required
optimizer	GradientTransformation	optax.GradientTransformation.	required
ensemble_size	int	Number of independent ensemble members E.	required
seed	PRNGKeyArray	PRNG key, split into E per-member init keys.	required

Returns:

Type	Description
EnsembleState	EnsembleState with stacked params and opt_state.
EnsembleState	Array leaves carry a leading (E,) axis.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def ensemble_init(
    init_fn: Callable[[PRNGKeyArray], PyTree],
    optimizer: optax.GradientTransformation,
    *,
    ensemble_size: int,
    seed: PRNGKeyArray,
) -> EnsembleState:
    """Initialize an ensemble of (params, opt_state) by vmap over keys.

    Args:
        init_fn: ``key -> params``. Called once per ensemble member.
        optimizer: ``optax.GradientTransformation``.
        ensemble_size: Number of independent ensemble members ``E``.
        seed: PRNG key, split into ``E`` per-member init keys.

    Returns:
        `EnsembleState` with stacked ``params`` and ``opt_state``.
        Array leaves carry a leading ``(E,)`` axis.
    """
    keys = jr.split(seed, ensemble_size)

    @eqx.filter_vmap
    def _per_member(key: PRNGKeyArray) -> tuple[PyTree, Any]:
        params = init_fn(key)
        opt_state = optimizer.init(eqx.filter(params, eqx.is_inexact_array))
        return params, opt_state

    params, opt_state = _per_member(keys)
    return EnsembleState(params=params, opt_state=opt_state)

ensemble_loss(log_joint: Callable[[PyTree, Array, Array], tuple[Float[Array, ''], Float[Array, '']]], *, prior_weight: float = 1.0, scale: float = 1.0) -> Callable[[PyTree, Array, Array], tuple[Float[Array, ''], PyTree]]

Build the filter_value_and_grad loss from a log-joint.

Returns a function loss_fn(params, x_batch, y_batch) -> (loss, grads) that computes

\[ \mathcal{L}(\theta) = -\, \text{scale} \cdot \log p(y \mid x, \theta) \;-\; w_{\text{prior}} \cdot \log p(\theta). \]

prior_weight=0 short-circuits the prior term so the user may return a placeholder 0.0 for logprior.

Parameters:

Name	Type	Description	Default
log_joint	Callable[[PyTree, Array, Array], tuple[Float[Array, ''], Float[Array, '']]]	(params, x_batch, y_batch) -> (loglik, logprior).	required
prior_weight	float	Weight on the logprior term. 0.0 ⇒ MLE, 1.0 ⇒ MAP.	1.0
scale	float	Multiplicative weight on the loglik term. Set to N / |B| for unbiased mini-batch SGD-MAP.	1.0

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def ensemble_loss(
    log_joint: Callable[
        [PyTree, Array, Array],
        tuple[Float[Array, ""], Float[Array, ""]],
    ],
    *,
    prior_weight: float = 1.0,
    scale: float = 1.0,
) -> Callable[[PyTree, Array, Array], tuple[Float[Array, ""], PyTree]]:
    r"""Build the ``filter_value_and_grad`` loss from a log-joint.

    Returns a function ``loss_fn(params, x_batch, y_batch) -> (loss, grads)``
    that computes

    $$
    \mathcal{L}(\theta) = -\, \text{scale} \cdot \log p(y \mid x, \theta)
        \;-\; w_{\text{prior}} \cdot \log p(\theta).
    $$

    ``prior_weight=0`` short-circuits the prior term so the user may
    return a placeholder ``0.0`` for ``logprior``.

    Args:
        log_joint: ``(params, x_batch, y_batch) -> (loglik, logprior)``.
        prior_weight: Weight on the ``logprior`` term. ``0.0`` ⇒ MLE,
            ``1.0`` ⇒ MAP.
        scale: Multiplicative weight on the ``loglik`` term. Set to
            ``N / |B|`` for unbiased mini-batch SGD-MAP.
    """
    use_prior = prior_weight != 0.0

    @eqx.filter_value_and_grad
    def _loss(params: PyTree, xb: Array, yb: Array) -> Float[Array, ""]:
        ll, lp = log_joint(params, xb, yb)
        nll = -scale * ll
        if use_prior:
            return nll - prior_weight * lp
        return nll

    return _loss

ensemble_step(state: EnsembleState, x_batch: Array, y_batch: Array, *, log_joint: Callable[[PyTree, Array, Array], tuple[Float[Array, ''], Float[Array, '']]], optimizer: optax.GradientTransformation, prior_weight: float = 1.0, scale: float = 1.0) -> tuple[EnsembleState, Float[Array, ' E']]

Perform one ensemble update step on a batch.

Each ensemble member computes ∇L(θ_e) independently (vmapped), advances its optax state, and returns the updated params + state.

Parameters:

Name	Type	Description	Default
state	EnsembleState	Current EnsembleState.	required
x_batch	Array	Inputs.	required
y_batch	Array	Targets.	required
log_joint	Callable[[PyTree, Array, Array], tuple[Float[Array, ''], Float[Array, '']]]	(params, x, y) -> (loglik, logprior).	required
optimizer	GradientTransformation	Same optax.GradientTransformation passed to ensemble_init. Stateless transforms can be re-built per call; stateful schedules need to share the same instance.	required
prior_weight	float	Weight on logprior; 0 ⇒ MLE.	1.0
scale	float	Weight on loglik (use N / |B| for mini-batch).	1.0

Returns:

Type	Description
EnsembleState	(new_state, per_member_losses) where losses has shape
Float[Array, ' E']	(E,).

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def ensemble_step(
    state: EnsembleState,
    x_batch: Array,
    y_batch: Array,
    *,
    log_joint: Callable[
        [PyTree, Array, Array],
        tuple[Float[Array, ""], Float[Array, ""]],
    ],
    optimizer: optax.GradientTransformation,
    prior_weight: float = 1.0,
    scale: float = 1.0,
) -> tuple[EnsembleState, Float[Array, " E"]]:
    """Perform one ensemble update step on a batch.

    Each ensemble member computes ``∇L(θ_e)`` independently (vmapped),
    advances its optax state, and returns the updated params + state.

    Args:
        state: Current `EnsembleState`.
        x_batch: Inputs.
        y_batch: Targets.
        log_joint: ``(params, x, y) -> (loglik, logprior)``.
        optimizer: Same ``optax.GradientTransformation`` passed to
            `ensemble_init`. Stateless transforms can be re-built
            per call; stateful schedules need to share the same
            instance.
        prior_weight: Weight on ``logprior``; ``0`` ⇒ MLE.
        scale: Weight on ``loglik`` (use ``N / |B|`` for mini-batch).

    Returns:
        ``(new_state, per_member_losses)`` where losses has shape
        ``(E,)``.
    """
    loss_fn = ensemble_loss(log_joint, prior_weight=prior_weight, scale=scale)

    @eqx.filter_vmap
    def _per_member(
        params: PyTree, opt_state: Any
    ) -> tuple[PyTree, Any, Float[Array, ""]]:
        loss, grads = loss_fn(params, x_batch, y_batch)
        updates, opt_state = optimizer.update(
            grads, opt_state, eqx.filter(params, eqx.is_inexact_array)
        )
        params = eqx.apply_updates(params, updates)
        return params, opt_state, loss

    params, opt_state, losses = _per_member(state.params, state.opt_state)
    return EnsembleState(params=params, opt_state=opt_state), losses

Layer 2 — NumPyro-like inference ops

init / update / run triplets that mirror numpyro.infer.SVI.

EnsembleMAP

Bases: Module

NumPyro-like ensemble MAP/MLE runner.

Mirrors numpyro.infer.SVI's init / update / run triplet, but every operation is ensembled by vmap over the leading (E,) axis.

Per-member objective is the tempered negative log-posterior

\[ \mathcal{L}_e(\theta_e) = -\frac{N}{|B|}\sum_{i \in B} \log p(y_i \mid x_i, \theta_e) \;-\; w_{\text{prior}} \cdot \log p(\theta_e), \]

where \(w_{\text{prior}}\) is prior_weight.

Tempering. Down-weighting the prior has the same argmax as tempering the likelihood by \(\beta = 1 / w_{\text{prior}}\), since

\[ \beta \mathcal{L} + \mathcal{P} = \beta \left( \mathcal{L} + \tfrac{1}{\beta} \mathcal{P} \right) \]

and rescaling an objective does not move its argmax. The two are not interchangeable in general: prior_weight scales the entire log prior — hyperpriors included — while a likelihood-side \(\beta\) (e.g. lfr_factor, which applies numpyro.handlers.scale to the likelihood site only) leaves every prior at unit weight. When the two coincide on which terms they reweight, their gradients differ by the uniform factor \(\beta\) — same argmax, but a scale an adaptive optimizer still responds to. This holds per minibatch as well: the tempered gradient is exactly \(\beta\) times the prior-weighted one, so minibatching introduces no further difference between the formulations.

Ensembling over seeds is worth more than usual for multi-modal objectives such as pyrox_gp.LatentFactorGPPrior's. Note what the spread does and does not mean: where the latent priors are identical the objective is rotation-invariant, and seed-to-seed spread along that gauge is an unidentifiable coordinate artifact, not uncertainty — averaging raw parameters across members is meaningless there. Ensemble gauge-invariant predictions instead, or align the members first. (With distinct latent kernels the rotation changes the objective, so no such flat manifold exists.)

Examples:

>>> runner = EnsembleMAP(
...     log_joint=log_joint,
...     init_fn=init_fn,
...     optimizer=optax.adam(5e-3),
...     ensemble_size=16,
... )
>>> # numpyro-style three-method API
>>> state = runner.init(jr.PRNGKey(0))
>>> for _ in range(2000):
...     state, losses = runner.update(state, x, y)
>>> # or one-shot
>>> result = runner.run(jr.PRNGKey(0), 2000, x, y)

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

class EnsembleMAP(eqx.Module):
    r"""NumPyro-like ensemble MAP/MLE runner.

    Mirrors `numpyro.infer.SVI`'s ``init`` / ``update`` / ``run``
    triplet, but every operation is ensembled by ``vmap`` over the
    leading ``(E,)`` axis.

    Per-member objective is the tempered negative log-posterior

    $$
    \mathcal{L}_e(\theta_e) = -\frac{N}{|B|}\sum_{i \in B}
        \log p(y_i \mid x_i, \theta_e)
        \;-\; w_{\text{prior}} \cdot \log p(\theta_e),
    $$

    where $w_{\text{prior}}$ is `prior_weight`.

    **Tempering.** Down-weighting the prior has the same argmax as
    tempering the likelihood by $\beta = 1 / w_{\text{prior}}$, since

    $$
    \beta \mathcal{L} + \mathcal{P}
    = \beta \left( \mathcal{L} + \tfrac{1}{\beta} \mathcal{P} \right)
    $$

    and rescaling an objective does not move its argmax. The two are
    **not** interchangeable in general: ``prior_weight`` scales the
    *entire* log prior — hyperpriors included — while a likelihood-side
    $\beta$ (e.g. [`lfr_factor`][pyrox_gp.lfr_factor], which applies
    ``numpyro.handlers.scale`` to the likelihood site only) leaves every
    prior at unit weight. When the two coincide on which terms they
    reweight, their gradients differ by the uniform factor $\beta$ —
    same argmax, but a scale an adaptive optimizer still responds to.
    This holds per minibatch as well: the tempered gradient is exactly
    $\beta$ times the prior-weighted one, so minibatching introduces no
    further difference between the formulations.

    Ensembling over seeds is worth more than usual for multi-modal
    objectives such as `pyrox_gp.LatentFactorGPPrior`'s. Note what the
    spread does and does not mean: where the latent priors are identical
    the objective is rotation-invariant, and seed-to-seed spread *along
    that gauge* is an unidentifiable coordinate artifact, not
    uncertainty — averaging raw parameters across members is meaningless
    there. Ensemble gauge-invariant **predictions** instead, or align the
    members first. (With distinct latent kernels the rotation changes the
    objective, so no such flat manifold exists.)

    Examples:
        >>> runner = EnsembleMAP(
        ...     log_joint=log_joint,
        ...     init_fn=init_fn,
        ...     optimizer=optax.adam(5e-3),
        ...     ensemble_size=16,
        ... )
        >>> # numpyro-style three-method API
        >>> state = runner.init(jr.PRNGKey(0))
        >>> for _ in range(2000):
        ...     state, losses = runner.update(state, x, y)
        >>> # or one-shot
        >>> result = runner.run(jr.PRNGKey(0), 2000, x, y)
    """

    log_joint: Callable[
        [PyTree, Array, Array],
        tuple[Float[Array, ""], Float[Array, ""]],
    ]
    init_fn: Callable[[PRNGKeyArray], PyTree]
    optimizer: Any  # optax.GradientTransformation, but optax may not be installed
    ensemble_size: int = eqx.field(static=True, default=16)
    prior_weight: float = eqx.field(static=True, default=1.0)

    def init(self, seed: PRNGKeyArray) -> EnsembleState:
        """Initialize the ensemble. Mirrors ``numpyro.infer.SVI.init``."""
        return ensemble_init(
            self.init_fn,
            self.optimizer,
            ensemble_size=self.ensemble_size,
            seed=seed,
        )

    def update(
        self,
        state: EnsembleState,
        x_batch: Array,
        y_batch: Array,
        *,
        scale: float = 1.0,
    ) -> tuple[EnsembleState, Float[Array, " E"]]:
        """One ensemble update step. Mirrors ``numpyro.infer.SVI.update``.

        Args:
            state: Current `EnsembleState`.
            x_batch: Batch inputs.
            y_batch: Batch targets.
            scale: ``N / |B|`` for mini-batch unbiased SGD-MAP. Defaults
                to ``1.0`` (full-batch).
        """
        return ensemble_step(
            state,
            x_batch,
            y_batch,
            log_joint=self.log_joint,
            optimizer=self.optimizer,
            prior_weight=self.prior_weight,
            scale=scale,
        )

    def run(
        self,
        seed: PRNGKeyArray,
        num_epochs: int,
        x: Array,
        y: Array,
        *,
        batch_size: int | None = None,
    ) -> EnsembleResult:
        """Fit the ensemble end-to-end. Mirrors ``numpyro.infer.SVI.run``.

        Internally drives `ensemble_step` via ``lax.scan`` for
        speed; equivalent to a hand-written Python loop over
        `update`.

        Args:
            seed: PRNG key used for both init and (when applicable)
                mini-batch index permutation.
            num_epochs: Number of optimizer steps per member.
            x: Inputs.
            y: Targets.
            batch_size: Optional mini-batch size. ``None`` ⇒ full-batch.

        Returns:
            `EnsembleResult` with stacked final params + loss
            history of shape ``(E, num_epochs)``.
        """
        n = y.shape[0]
        # Clamp batch_size to n so that batch_size > n falls back cleanly
        # to full-batch with scale=1, instead of silently downweighting
        # the likelihood by n/batch_size < 1 and biasing toward the prior.
        bsz = n if batch_size is None else min(batch_size, n)
        scale = float(n) / float(bsz)
        use_minibatch = batch_size is not None and bsz < n

        init_key, perm_key = jr.split(seed)
        state = self.init(init_key)

        # Partition non-array leaves out of the scan carry; lax.scan only
        # accepts JAX-typed carry, but eqx.Module trees may contain
        # captured Python callables (e.g. jax.nn.tanh). Static is shared
        # across ensemble members, so we recombine after the scan.
        arrays0, static = eqx.partition(state.params, eqx.is_inexact_array)

        def epoch(
            carry: tuple[PyTree, Any, PRNGKeyArray],
            _: None,
        ) -> tuple[tuple[PyTree, Any, PRNGKeyArray], Float[Array, " E"]]:
            arrays, opt_state, key = carry
            if use_minibatch:
                key, sub = jr.split(key)
                idx = jr.permutation(sub, n)[:bsz]
                xb, yb = x[idx], y[idx]
            else:
                xb, yb = x, y
            params = eqx.combine(arrays, static)
            new_state, losses = ensemble_step(
                EnsembleState(params=params, opt_state=opt_state),
                xb,
                yb,
                log_joint=self.log_joint,
                optimizer=self.optimizer,
                prior_weight=self.prior_weight,
                scale=scale,
            )
            new_arrays, _ = eqx.partition(new_state.params, eqx.is_inexact_array)
            return (new_arrays, new_state.opt_state, key), losses

        (arrays_final, _, _), losses = jax.lax.scan(
            epoch, (arrays0, state.opt_state, perm_key), None, length=num_epochs
        )
        final_params = eqx.combine(arrays_final, static)
        # losses is (T, E) from scan; transpose to (E, T) per the public contract.
        return EnsembleResult(params=final_params, losses=einx.id("t e -> e t", losses))

init(seed: PRNGKeyArray) -> EnsembleState

Initialize the ensemble. Mirrors numpyro.infer.SVI.init.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def init(self, seed: PRNGKeyArray) -> EnsembleState:
    """Initialize the ensemble. Mirrors ``numpyro.infer.SVI.init``."""
    return ensemble_init(
        self.init_fn,
        self.optimizer,
        ensemble_size=self.ensemble_size,
        seed=seed,
    )

update(state: EnsembleState, x_batch: Array, y_batch: Array, *, scale: float = 1.0) -> tuple[EnsembleState, Float[Array, ' E']]

One ensemble update step. Mirrors numpyro.infer.SVI.update.

Parameters:

Name	Type	Description	Default
state	EnsembleState	Current EnsembleState.	required
x_batch	Array	Batch inputs.	required
y_batch	Array	Batch targets.	required
scale	float	N / |B| for mini-batch unbiased SGD-MAP. Defaults to 1.0 (full-batch).	1.0

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def update(
    self,
    state: EnsembleState,
    x_batch: Array,
    y_batch: Array,
    *,
    scale: float = 1.0,
) -> tuple[EnsembleState, Float[Array, " E"]]:
    """One ensemble update step. Mirrors ``numpyro.infer.SVI.update``.

    Args:
        state: Current `EnsembleState`.
        x_batch: Batch inputs.
        y_batch: Batch targets.
        scale: ``N / |B|`` for mini-batch unbiased SGD-MAP. Defaults
            to ``1.0`` (full-batch).
    """
    return ensemble_step(
        state,
        x_batch,
        y_batch,
        log_joint=self.log_joint,
        optimizer=self.optimizer,
        prior_weight=self.prior_weight,
        scale=scale,
    )

run(seed: PRNGKeyArray, num_epochs: int, x: Array, y: Array, *, batch_size: int | None = None) -> EnsembleResult

Fit the ensemble end-to-end. Mirrors numpyro.infer.SVI.run.

Internally drives ensemble_step via lax.scan for speed; equivalent to a hand-written Python loop over update.

Parameters:

Name	Type	Description	Default
seed	PRNGKeyArray	PRNG key used for both init and (when applicable) mini-batch index permutation.	required
num_epochs	int	Number of optimizer steps per member.	required
x	Array	Inputs.	required
y	Array	Targets.	required
batch_size	int | None	Optional mini-batch size. None ⇒ full-batch.	None

Returns:

Type	Description
EnsembleResult	EnsembleResult with stacked final params + loss
EnsembleResult	history of shape (E, num_epochs).

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def run(
    self,
    seed: PRNGKeyArray,
    num_epochs: int,
    x: Array,
    y: Array,
    *,
    batch_size: int | None = None,
) -> EnsembleResult:
    """Fit the ensemble end-to-end. Mirrors ``numpyro.infer.SVI.run``.

    Internally drives `ensemble_step` via ``lax.scan`` for
    speed; equivalent to a hand-written Python loop over
    `update`.

    Args:
        seed: PRNG key used for both init and (when applicable)
            mini-batch index permutation.
        num_epochs: Number of optimizer steps per member.
        x: Inputs.
        y: Targets.
        batch_size: Optional mini-batch size. ``None`` ⇒ full-batch.

    Returns:
        `EnsembleResult` with stacked final params + loss
        history of shape ``(E, num_epochs)``.
    """
    n = y.shape[0]
    # Clamp batch_size to n so that batch_size > n falls back cleanly
    # to full-batch with scale=1, instead of silently downweighting
    # the likelihood by n/batch_size < 1 and biasing toward the prior.
    bsz = n if batch_size is None else min(batch_size, n)
    scale = float(n) / float(bsz)
    use_minibatch = batch_size is not None and bsz < n

    init_key, perm_key = jr.split(seed)
    state = self.init(init_key)

    # Partition non-array leaves out of the scan carry; lax.scan only
    # accepts JAX-typed carry, but eqx.Module trees may contain
    # captured Python callables (e.g. jax.nn.tanh). Static is shared
    # across ensemble members, so we recombine after the scan.
    arrays0, static = eqx.partition(state.params, eqx.is_inexact_array)

    def epoch(
        carry: tuple[PyTree, Any, PRNGKeyArray],
        _: None,
    ) -> tuple[tuple[PyTree, Any, PRNGKeyArray], Float[Array, " E"]]:
        arrays, opt_state, key = carry
        if use_minibatch:
            key, sub = jr.split(key)
            idx = jr.permutation(sub, n)[:bsz]
            xb, yb = x[idx], y[idx]
        else:
            xb, yb = x, y
        params = eqx.combine(arrays, static)
        new_state, losses = ensemble_step(
            EnsembleState(params=params, opt_state=opt_state),
            xb,
            yb,
            log_joint=self.log_joint,
            optimizer=self.optimizer,
            prior_weight=self.prior_weight,
            scale=scale,
        )
        new_arrays, _ = eqx.partition(new_state.params, eqx.is_inexact_array)
        return (new_arrays, new_state.opt_state, key), losses

    (arrays_final, _, _), losses = jax.lax.scan(
        epoch, (arrays0, state.opt_state, perm_key), None, length=num_epochs
    )
    final_params = eqx.combine(arrays_final, static)
    # losses is (T, E) from scan; transpose to (E, T) per the public contract.
    return EnsembleResult(params=final_params, losses=einx.id("t e -> e t", losses))

EnsembleVI

Bases: Module

NumPyro-like ensemble variational-inference runner.

Wraps numpyro.infer.SVI + numpyro.infer.Trace_ELBO with the same ensemble surface as EnsembleMAP.

Per-member objective is the tempered ELBO

\[ \mathrm{ELBO}(\phi_e) = \mathbb{E}_{q_{\phi_e}}\! \bigl[\log p(y \mid x, \theta)\bigr] - \beta\, \mathrm{KL}\!\bigl(q_{\phi_e}\,\|\,p\bigr), \]

where \(\beta\) is kl_weight.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

class EnsembleVI(eqx.Module):
    r"""NumPyro-like ensemble variational-inference runner.

    Wraps `numpyro.infer.SVI` + `numpyro.infer.Trace_ELBO`
    with the same ensemble surface as `EnsembleMAP`.

    Per-member objective is the tempered ELBO

    $$
    \mathrm{ELBO}(\phi_e) = \mathbb{E}_{q_{\phi_e}}\!
        \bigl[\log p(y \mid x, \theta)\bigr]
        - \beta\, \mathrm{KL}\!\bigl(q_{\phi_e}\,\|\,p\bigr),
    $$

    where $\beta$ is `kl_weight`.
    """

    model_fn: Callable[..., None]
    guide_fn: Callable[..., None]
    optimizer: Any  # numpyro.optim or optax.GradientTransformation
    ensemble_size: int = eqx.field(static=True, default=16)
    kl_weight: float = eqx.field(static=True, default=1.0)
    num_particles: int = eqx.field(static=True, default=1)

    def _build_svi(self) -> Any:
        import numpyro
        from numpyro.infer import SVI, Trace_ELBO

        opt = self.optimizer
        # Duck-type rather than isinstance against numpyro.optim._NumPyroOptim:
        # the latter is a private symbol that may move under refactors. The
        # public contract of a numpyro optimizer is `init_fn` + `update_fn` +
        # `get_params`; if any are missing, assume it's an optax transform.
        is_numpyro_optim = (
            callable(getattr(opt, "init_fn", None))
            and callable(getattr(opt, "update_fn", None))
            and callable(getattr(opt, "get_params", None))
        )
        if not is_numpyro_optim:
            opt = numpyro.optim.optax_to_numpyro(opt)
        if self.kl_weight == 1.0:
            elbo: Any = Trace_ELBO(num_particles=self.num_particles)
        else:
            elbo = _TemperedTraceELBO(
                kl_weight=self.kl_weight, num_particles=self.num_particles
            )
        return SVI(self.model_fn, self.guide_fn, opt, loss=elbo)

    def init(self, seed: PRNGKeyArray, *args: Any, **kwargs: Any) -> Any:
        """Initialize the ensemble of SVI states."""
        svi = self._build_svi()
        keys = jr.split(seed, self.ensemble_size)
        return jax.vmap(lambda k: svi.init(k, *args, **kwargs))(keys)

    def update(
        self, state: Any, *args: Any, **kwargs: Any
    ) -> tuple[Any, Float[Array, " E"]]:
        """One ensemble SVI update. Mirrors ``numpyro.infer.SVI.update``."""
        svi = self._build_svi()
        return jax.vmap(lambda s: svi.update(s, *args, **kwargs))(state)

    def run(
        self,
        seed: PRNGKeyArray,
        num_epochs: int,
        *args: Any,
        **kwargs: Any,
    ) -> EnsembleResult:
        """Fit the ensemble end-to-end via vmapped ``svi.run``."""
        svi = self._build_svi()
        keys = jr.split(seed, self.ensemble_size)

        @jax.vmap
        def _run_one(
            k: PRNGKeyArray,
        ) -> tuple[PyTree, Float[Array, " T"]]:
            r = svi.run(k, num_epochs, *args, progress_bar=False, **kwargs)
            return r.params, r.losses

        params, losses = _run_one(keys)
        return EnsembleResult(params=params, losses=losses)

init(seed: PRNGKeyArray, *args: Any, **kwargs: Any) -> Any

Initialize the ensemble of SVI states.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def init(self, seed: PRNGKeyArray, *args: Any, **kwargs: Any) -> Any:
    """Initialize the ensemble of SVI states."""
    svi = self._build_svi()
    keys = jr.split(seed, self.ensemble_size)
    return jax.vmap(lambda k: svi.init(k, *args, **kwargs))(keys)

update(state: Any, *args: Any, **kwargs: Any) -> tuple[Any, Float[Array, ' E']]

One ensemble SVI update. Mirrors numpyro.infer.SVI.update.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def update(
    self, state: Any, *args: Any, **kwargs: Any
) -> tuple[Any, Float[Array, " E"]]:
    """One ensemble SVI update. Mirrors ``numpyro.infer.SVI.update``."""
    svi = self._build_svi()
    return jax.vmap(lambda s: svi.update(s, *args, **kwargs))(state)

run(seed: PRNGKeyArray, num_epochs: int, *args: Any, **kwargs: Any) -> EnsembleResult

Fit the ensemble end-to-end via vmapped svi.run.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def run(
    self,
    seed: PRNGKeyArray,
    num_epochs: int,
    *args: Any,
    **kwargs: Any,
) -> EnsembleResult:
    """Fit the ensemble end-to-end via vmapped ``svi.run``."""
    svi = self._build_svi()
    keys = jr.split(seed, self.ensemble_size)

    @jax.vmap
    def _run_one(
        k: PRNGKeyArray,
    ) -> tuple[PyTree, Float[Array, " T"]]:
        r = svi.run(k, num_epochs, *args, progress_bar=False, **kwargs)
        return r.params, r.losses

    params, losses = _run_one(keys)
    return EnsembleResult(params=params, losses=losses)

Layer 3 — One-shot sugar

ensemble_map(log_joint: Callable[[PyTree, Array, Array], tuple[Float[Array, ''], Float[Array, '']]], init_fn: Callable[[PRNGKeyArray], PyTree], *, ensemble_size: int, num_epochs: int, data: tuple[Array, Array], seed: PRNGKeyArray, batch_size: int | None = None, learning_rate: float = 0.005, prior_weight: float = 1.0, optimizer: optax.GradientTransformation | None = None) -> tuple[PyTree, Float[Array, 'E T']]

One-shot wrapper around EnsembleMAP.

Equivalent to EnsembleMAP(log_joint, init_fn, optimizer, ensemble_size=E, prior_weight=w).run(seed, num_epochs, *data, batch_size=B).

Parameters:

Name	Type	Description	Default
log_joint	Callable[[PyTree, Array, Array], tuple[Float[Array, ''], Float[Array, '']]]	(params, x_batch, y_batch) -> (loglik, logprior).	required
init_fn	Callable[[PRNGKeyArray], PyTree]	key -> params.	required
ensemble_size	int	Number of independent MAP fits E.	required
num_epochs	int	Optimizer steps per member.	required
data	tuple[Array, Array]	(x, y).	required
seed	PRNGKeyArray	PRNG key.	required
batch_size	int | None	Mini-batch size. None ⇒ full-batch.	None
learning_rate	float	Default-Adam learning rate. Ignored if optimizer is supplied.	0.005
prior_weight	float	0 ⇒ MLE, 1 ⇒ MAP.	1.0
optimizer	GradientTransformation | None	Optional optax.GradientTransformation. Defaults to optax.adam(learning_rate).	None

Returns:

Type	Description
PyTree	(params_stacked, losses) — leading (E,) axis on
Float[Array, 'E T']	params; losses shape (E, num_epochs).

Examples:

>>> params, losses = ensemble_map(
...     log_joint, init_fn,
...     ensemble_size=16, num_epochs=2000,
...     data=(X, y), seed=jr.PRNGKey(0),
... )

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def ensemble_map(
    log_joint: Callable[
        [PyTree, Array, Array],
        tuple[Float[Array, ""], Float[Array, ""]],
    ],
    init_fn: Callable[[PRNGKeyArray], PyTree],
    *,
    ensemble_size: int,
    num_epochs: int,
    data: tuple[Array, Array],
    seed: PRNGKeyArray,
    batch_size: int | None = None,
    learning_rate: float = 5e-3,
    prior_weight: float = 1.0,
    optimizer: optax.GradientTransformation | None = None,
) -> tuple[PyTree, Float[Array, "E T"]]:
    """One-shot wrapper around `EnsembleMAP`.

    Equivalent to ``EnsembleMAP(log_joint, init_fn, optimizer,
    ensemble_size=E, prior_weight=w).run(seed, num_epochs, *data,
    batch_size=B)``.

    Args:
        log_joint: ``(params, x_batch, y_batch) -> (loglik, logprior)``.
        init_fn: ``key -> params``.
        ensemble_size: Number of independent MAP fits ``E``.
        num_epochs: Optimizer steps per member.
        data: ``(x, y)``.
        seed: PRNG key.
        batch_size: Mini-batch size. ``None`` ⇒ full-batch.
        learning_rate: Default-Adam learning rate. Ignored if
            ``optimizer`` is supplied.
        prior_weight: ``0`` ⇒ MLE, ``1`` ⇒ MAP.
        optimizer: Optional ``optax.GradientTransformation``. Defaults
            to ``optax.adam(learning_rate)``.

    Returns:
        ``(params_stacked, losses)`` — leading ``(E,)`` axis on
        ``params``; ``losses`` shape ``(E, num_epochs)``.

    Examples:
        >>> params, losses = ensemble_map(
        ...     log_joint, init_fn,
        ...     ensemble_size=16, num_epochs=2000,
        ...     data=(X, y), seed=jr.PRNGKey(0),
        ... )
    """
    optax = _require_optax()
    opt = optimizer if optimizer is not None else optax.adam(learning_rate)
    runner = EnsembleMAP(
        log_joint=log_joint,
        init_fn=init_fn,
        optimizer=opt,
        ensemble_size=ensemble_size,
        prior_weight=prior_weight,
    )
    result = runner.run(seed, num_epochs, *data, batch_size=batch_size)
    return result.params, result.losses

ensemble_vi(model_fn: Callable[..., None], guide_fn: Callable[..., None], *, ensemble_size: int, num_epochs: int, data: tuple[Array, Array], seed: PRNGKeyArray, kl_weight: float = 1.0, learning_rate: float = 0.005, optimizer: Any = None, num_particles: int = 1) -> tuple[PyTree, Float[Array, 'E T']]

One-shot wrapper around EnsembleVI.

Parameters:

Name	Type	Description	Default
model_fn	Callable[..., None]	NumPyro model (x, y) -> None.	required
guide_fn	Callable[..., None]	NumPyro guide.	required
ensemble_size	int	Number of SVI fits E.	required
num_epochs	int	Steps per member.	required
data	tuple[Array, Array]	(x, y).	required
seed	PRNGKeyArray	PRNG key.	required
kl_weight	float	ELBO temper β. 1.0 ⇒ standard ELBO.	1.0
learning_rate	float	Default-Adam learning rate.	0.005
optimizer	Any	Optional numpyro.optim or optax.GradientTransformation (auto-wrapped).	None
num_particles	int	MC particles per ELBO estimate.	1

Returns:

Type	Description
tuple[PyTree, Float[Array, 'E T']]	(guide_params_stacked, losses) with leading (E,) axis.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def ensemble_vi(
    model_fn: Callable[..., None],
    guide_fn: Callable[..., None],
    *,
    ensemble_size: int,
    num_epochs: int,
    data: tuple[Array, Array],
    seed: PRNGKeyArray,
    kl_weight: float = 1.0,
    learning_rate: float = 5e-3,
    optimizer: Any = None,
    num_particles: int = 1,
) -> tuple[PyTree, Float[Array, "E T"]]:
    """One-shot wrapper around `EnsembleVI`.

    Args:
        model_fn: NumPyro model ``(x, y) -> None``.
        guide_fn: NumPyro guide.
        ensemble_size: Number of SVI fits ``E``.
        num_epochs: Steps per member.
        data: ``(x, y)``.
        seed: PRNG key.
        kl_weight: ELBO temper ``β``. ``1.0`` ⇒ standard ELBO.
        learning_rate: Default-Adam learning rate.
        optimizer: Optional ``numpyro.optim`` or
            ``optax.GradientTransformation`` (auto-wrapped).
        num_particles: MC particles per ELBO estimate.

    Returns:
        ``(guide_params_stacked, losses)`` with leading ``(E,)`` axis.
    """
    import numpyro

    if optimizer is None:
        optimizer = numpyro.optim.Adam(learning_rate)
    runner = EnsembleVI(
        model_fn=model_fn,
        guide_fn=guide_fn,
        optimizer=optimizer,
        ensemble_size=ensemble_size,
        kl_weight=kl_weight,
        num_particles=num_particles,
    )
    result = runner.run(seed, num_epochs, *data)
    return result.params, result.losses

ensemble_predict(params_stacked: PyTree, predict_fn: Callable[[PyTree, Array], Array], x_new: Array) -> Array

Vmap predict_fn over the leading ensemble axis of params.

Uses equinox.filter_vmap so it works whether params_stacked is a pure-array PyTree or an equinox.Module containing non-array leaves (e.g. captured jax.nn.tanh). Array leaves are mapped over axis 0; non-array leaves are broadcast.

Parameters:

Name	Type	Description	Default
params_stacked	PyTree	PyTree returned by ensemble_map / ensemble_vi / EnsembleMAP.run; every array leaf has a leading (E,) axis.	required
predict_fn	Callable[[PyTree, Array], Array]	(params, x) -> y.	required
x_new	Array	Inputs to predict at; shared across all members.	required

Returns:

Type	Description
Array	Stacked predictions with leading (E,) axis.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

def ensemble_predict(
    params_stacked: PyTree,
    predict_fn: Callable[[PyTree, Array], Array],
    x_new: Array,
) -> Array:
    """Vmap ``predict_fn`` over the leading ensemble axis of params.

    Uses `equinox.filter_vmap` so it works whether
    ``params_stacked`` is a pure-array PyTree or an
    `equinox.Module` containing non-array leaves (e.g. captured
    ``jax.nn.tanh``). Array leaves are mapped over axis 0; non-array
    leaves are broadcast.

    Args:
        params_stacked: PyTree returned by `ensemble_map` /
            `ensemble_vi` / `EnsembleMAP.run`; every array
            leaf has a leading ``(E,)`` axis.
        predict_fn: ``(params, x) -> y``.
        x_new: Inputs to predict at; shared across all members.

    Returns:
        Stacked predictions with leading ``(E,)`` axis.
    """
    return eqx.filter_vmap(predict_fn, in_axes=(eqx.if_array(0), None))(
        params_stacked, x_new
    )

Post-hoc Laplace

A Gaussian posterior over a chosen parameter group (typically the last layer) of one fitted model, from the generalized Gauss-Newton Hessian, with a linearized predictive and evidence-tuned prior precision. Needs pip install pyrox[laplace].

last_layer_laplace(predict_fn: Callable[[PyTree, Array], Array], params: PyTree, x: Array, y: Array, *, sigma: Float[Array, ''] | float, param_filter: Callable[[tuple, Any], bool], prior_precision: Float[Array, ''] | float = 1.0, structure: Literal['full', 'diag'] = 'full', jitter: float = 1e-08) -> LastLayerLaplace

Fit a post-hoc GGN Laplace posterior over a parameter group.

The fitted model is represented as in ensemble_map / ensemble_predict: a parameter PyTree (one member, no ensemble axis; an equinox.Module works) and predict_fn(params, x) -> f. The weights \(w\) selected by param_filter get a Gaussian posterior; every other leaf stays at its MAP value. With \(J = \partial f / \partial w\) on the training inputs, \(R = \sigma^2 I\) and prior \(\mathcal{N}(0, \lambda^{-1} I)\),

\[ \Sigma_w^{-1} = J^\top R^{-1} J + \lambda I, \qquad \mathrm{Var}[f(x_*)] = J_* \Sigma_w J_*^\top. \]

For a last layer that is linear in \(w\), \(J\) is the frozen feature matrix and this is exact Bayesian linear regression on those features. \(J\) is formed densely with forward-mode autodiff, one pass per weight, so keep the group small (a last layer, not the network).

Parameters:

Name	Type	Description	Default
predict_fn	Callable[[PyTree, Array], Array]	(params, x) -> f; any output shape (flattened for the Jacobian).	required
params	PyTree	MAP parameters, e.g. one member of ensemble_map's output.	required
x	Array	Training inputs.	required
y	Array	Training targets, shaped like predict_fn(params, x).	required
sigma	Float[Array, ''] | float	Gaussian observation noise standard deviation.	required
param_filter	Callable[[tuple, Any], bool]	(path, leaf) -> bool on jax.tree_util key paths, as param_group_optimizer's label_fn; only inexact array leaves are considered.	required
prior_precision	Float[Array, ''] | float	Isotropic prior precision \(\lambda\).	1.0
structure	Literal['full', 'diag']	"full" keeps \(J^\top J\); "diag" keeps its diagonal (gaussx.ggn_diagonal).	'full'
jitter	float	Added to the diagonal of the GGN (so of the precision) before factorizing, consistently in every use.	1e-08

Returns:

Type	Description
LastLayerLaplace	A LastLayerLaplace posterior.

Raises:

Type	Description
ValueError	If structure is unknown or param_filter selects no weights.

Examples:

>>> import equinox as eqx
>>> import jax
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> from pyrox.inference import last_layer_laplace
>>> model = eqx.nn.MLP(1, 1, 16, 2, activation=jnp.sin, key=jr.key(0))
>>> X = jnp.linspace(-1.0, 1.0, 50)[:, None]
>>> y = jnp.sin(3.0 * X[:, 0])
>>> def predict(m, x):
...     return jax.vmap(m)(x)[:, 0]
>>> # Key paths spell list indices out (".layers[2]"), so build the
>>> # selector from the index rather than writing "layers[-1]".
>>> last = f".layers[{len(model.layers) - 1}]"
>>> def is_last(path, _leaf):
...     return jax.tree_util.keystr(path).startswith(last)
>>> post = last_layer_laplace(
...     predict, model, X, y, sigma=0.1, param_filter=is_last
... )
>>> post.w_map.shape  # 16 weights + 1 bias
(17,)
>>> post = post.optimize_prior_precision()
>>> mean, var = post.predict(jnp.linspace(-1.5, 1.5, 7)[:, None])
>>> mean.shape, var.shape
((7,), (7,))

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

def last_layer_laplace(
    predict_fn: Callable[[PyTree, Array], Array],
    params: PyTree,
    x: Array,
    y: Array,
    *,
    sigma: Float[Array, ""] | float,
    param_filter: Callable[[tuple, Any], bool],
    prior_precision: Float[Array, ""] | float = 1.0,
    structure: Literal["full", "diag"] = "full",
    jitter: float = 1e-8,
) -> LastLayerLaplace:
    r"""Fit a post-hoc GGN Laplace posterior over a parameter group.

    The fitted model is represented as in `ensemble_map` /
    `ensemble_predict`: a parameter PyTree (one member, no ensemble axis;
    an `equinox.Module` works) and ``predict_fn(params, x) -> f``. The
    weights $w$ selected by ``param_filter`` get a Gaussian posterior; every
    other leaf stays at its MAP value. With $J = \partial f / \partial w$ on
    the training inputs, $R = \sigma^2 I$ and prior $\mathcal{N}(0,
    \lambda^{-1} I)$,

    $$
    \Sigma_w^{-1} = J^\top R^{-1} J + \lambda I, \qquad
    \mathrm{Var}[f(x_*)] = J_* \Sigma_w J_*^\top.
    $$

    For a last layer that is linear in $w$, $J$ is the frozen feature
    matrix and this is exact Bayesian linear regression on those features.
    $J$ is formed densely with forward-mode autodiff, one pass per weight,
    so keep the group small (a last layer, not the network).

    Args:
        predict_fn: ``(params, x) -> f``; any output shape (flattened for
            the Jacobian).
        params: MAP parameters, e.g. one member of `ensemble_map`'s output.
        x: Training inputs.
        y: Training targets, shaped like ``predict_fn(params, x)``.
        sigma: Gaussian observation noise standard deviation.
        param_filter: ``(path, leaf) -> bool`` on ``jax.tree_util`` key
            paths, as `param_group_optimizer`'s ``label_fn``; only inexact
            array leaves are considered.
        prior_precision: Isotropic prior precision $\lambda$.
        structure: ``"full"`` keeps $J^\top J$; ``"diag"`` keeps its
            diagonal (``gaussx.ggn_diagonal``).
        jitter: Added to the diagonal of the GGN (so of the precision)
            before factorizing, consistently in every use.

    Returns:
        A `LastLayerLaplace` posterior.

    Raises:
        ValueError: If ``structure`` is unknown or ``param_filter`` selects
            no weights.

    Examples:
        >>> import equinox as eqx
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import jax.random as jr
        >>> from pyrox.inference import last_layer_laplace
        >>> model = eqx.nn.MLP(1, 1, 16, 2, activation=jnp.sin, key=jr.key(0))
        >>> X = jnp.linspace(-1.0, 1.0, 50)[:, None]
        >>> y = jnp.sin(3.0 * X[:, 0])
        >>> def predict(m, x):
        ...     return jax.vmap(m)(x)[:, 0]
        >>> # Key paths spell list indices out (".layers[2]"), so build the
        >>> # selector from the index rather than writing "layers[-1]".
        >>> last = f".layers[{len(model.layers) - 1}]"
        >>> def is_last(path, _leaf):
        ...     return jax.tree_util.keystr(path).startswith(last)
        >>> post = last_layer_laplace(
        ...     predict, model, X, y, sigma=0.1, param_filter=is_last
        ... )
        >>> post.w_map.shape  # 16 weights + 1 bias
        (17,)
        >>> post = post.optimize_prior_precision()
        >>> mean, var = post.predict(jnp.linspace(-1.5, 1.5, 7)[:, None])
        >>> mean.shape, var.shape
        ((7,), (7,))
    """
    gaussx, _ = _require_gaussx()
    if structure not in ("full", "diag"):
        raise ValueError(f"structure must be 'full' or 'diag', got {structure!r}.")
    w, unravel, rest = _split(params, param_filter)
    if w.shape[0] == 0:
        raise ValueError("param_filter selected no inexact array leaves of params.")
    f = _forward(predict_fn, unravel, rest, x)
    sigma = jnp.asarray(sigma, dtype=w.dtype)
    jac = jax.jacfwd(f)(w)  # (N, P)
    f_map = f(w)
    if structure == "diag":
        ggn = gaussx.ggn_diagonal(jac) / sigma**2
    else:
        # The dense P x P GGN, not gaussx.gauss_newton_precision's Woodbury
        # form: optimize_prior_precision eigendecomposes G once to search
        # lambda in O(P) per step, and storing G (P^2) instead of J (N P)
        # is what lets N, the number of field samples, far exceed P.
        ggn = einx.dot("n p, n q -> p q", jac, jac) / sigma**2
    resid = jnp.ravel(y) - f_map
    loglik_grad = einx.dot("n p, n -> p", jac, resid) / sigma**2
    log_lik = dist.Normal(f_map, sigma).log_prob(jnp.ravel(y)).sum()
    return LastLayerLaplace(
        predict_fn=predict_fn,
        params=params,
        param_filter=param_filter,
        w_map=w,
        w_mean=w,
        ggn=ggn,
        loglik_grad=loglik_grad,
        log_likelihood=log_lik,
        sigma=sigma,
        prior_precision=jnp.asarray(prior_precision, dtype=w.dtype),
        structure=structure,
        jitter=jitter,
    )

LastLayerLaplace

Bases: Module

Gaussian posterior over a selected parameter group of a fitted model.

Built by last_layer_laplace. With \(w\) the selected weights (the rest frozen at the MAP), \(J\) the Jacobian \(\partial f / \partial w\) at the linearization point \(w^*\) on the training inputs and Gaussian noise \(\sigma^2\), the log-likelihood is replaced by its Gauss-Newton quadratic model around \(w^*\),

\[ \log p(y \mid w) \approx \ell^* + b^\top \Delta - \tfrac{1}{2} \Delta^\top G \Delta, \qquad \Delta = w - w^*,\; G = \sigma^{-2} J^\top J + \epsilon I,\; b = \sigma^{-2} J^\top (y - f^*), \]

with \(\epsilon\) the jitter, which is exact for a model linear in \(w\). Under the prior \(\mathcal{N}(0, \lambda^{-1} I)\) the posterior precision is \(H_\lambda = G + \lambda I\) and its mode is \(w_\lambda = H_\lambda^{-1} (G w^* + b)\). \(G\) is kept whole (structure="full") or as its diagonal (structure="diag"). The posterior registers no NumPyro sites.

Attributes:

Name	Type	Description
predict_fn	Callable[[PyTree, Array], Array]	(params, x) -> f, the same predictive as ensemble_predict takes.
params	PyTree	The fitted parameters; their selected weights are \(w^*\).
param_filter	Callable[[tuple, Any], bool]	(path, leaf) -> bool selecting the weights \(w\).
w_map	Float[Array, ' P']	The linearization point \(w^*\), flattened, shape (P,).
w_mean	Float[Array, ' P']	The posterior mean, shape (P,): \(w^*\) as fitted, or \(w_\lambda\) after optimize_prior_precision.
ggn	Float[Array, 'P P'] | Float[Array, ' P']	\(\sigma^{-2} J^\top J\) (without jitter), shape (P, P) (full) or (P,) (diag).
loglik_grad	Float[Array, ' P']	\(b\), the log-likelihood gradient at \(w^*\), (P,).
log_likelihood	Float[Array, '']	\(\ell^* = \log p(y \mid w^*)\) on the training data.
sigma	Float[Array, '']	Observation noise standard deviation \(\sigma\).
prior_precision	Float[Array, '']	Isotropic prior precision \(\lambda\).
structure	Literal['full', 'diag']	"full" or "diag".
jitter	float	\(\epsilon\), added to the diagonal of \(G\), and so of \(H\), everywhere it is used.

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

class LastLayerLaplace(eqx.Module):
    r"""Gaussian posterior over a selected parameter group of a fitted model.

    Built by [`last_layer_laplace`][pyrox.inference.last_layer_laplace].
    With $w$ the selected weights (the rest frozen at the MAP), $J$ the
    Jacobian $\partial f / \partial w$ at the linearization point $w^*$ on
    the training inputs and Gaussian noise $\sigma^2$, the log-likelihood is
    replaced by its Gauss-Newton quadratic model around $w^*$,

    $$
    \log p(y \mid w) \approx \ell^* + b^\top \Delta
        - \tfrac{1}{2} \Delta^\top G \Delta, \qquad
    \Delta = w - w^*,\;
    G = \sigma^{-2} J^\top J + \epsilon I,\;
    b = \sigma^{-2} J^\top (y - f^*),
    $$

    with $\epsilon$ the ``jitter``, which is exact for a model linear in
    $w$. Under the prior $\mathcal{N}(0, \lambda^{-1} I)$ the posterior
    precision is $H_\lambda = G + \lambda I$ and its mode is
    $w_\lambda = H_\lambda^{-1} (G w^* + b)$. $G$ is kept whole
    (``structure="full"``) or as its diagonal (``structure="diag"``). The
    posterior registers no NumPyro sites.

    Attributes:
        predict_fn: ``(params, x) -> f``, the same predictive as
            `ensemble_predict` takes.
        params: The fitted parameters; their selected weights are $w^*$.
        param_filter: ``(path, leaf) -> bool`` selecting the weights $w$.
        w_map: The linearization point $w^*$, flattened, shape ``(P,)``.
        w_mean: The posterior mean, shape ``(P,)``: $w^*$ as fitted, or
            $w_\lambda$ after `optimize_prior_precision`.
        ggn: $\sigma^{-2} J^\top J$ (without jitter), shape ``(P, P)``
            (full) or ``(P,)`` (diag).
        loglik_grad: $b$, the log-likelihood gradient at $w^*$, ``(P,)``.
        log_likelihood: $\ell^* = \log p(y \mid w^*)$ on the training data.
        sigma: Observation noise standard deviation $\sigma$.
        prior_precision: Isotropic prior precision $\lambda$.
        structure: ``"full"`` or ``"diag"``.
        jitter: $\epsilon$, added to the diagonal of $G$, and so of $H$,
            everywhere it is used.
    """

    predict_fn: Callable[[PyTree, Array], Array] = eqx.field(static=True)
    params: PyTree
    param_filter: Callable[[tuple, Any], bool] = eqx.field(static=True)
    w_map: Float[Array, " P"]
    w_mean: Float[Array, " P"]
    ggn: Float[Array, "P P"] | Float[Array, " P"]
    loglik_grad: Float[Array, " P"]
    log_likelihood: Float[Array, ""]
    sigma: Float[Array, ""]
    prior_precision: Float[Array, ""]
    structure: Literal["full", "diag"] = eqx.field(static=True)
    jitter: float = eqx.field(static=True)

    def _lam(self, prior_precision: Array | float | None) -> Array:
        if prior_precision is None:
            return self.prior_precision
        return jnp.asarray(prior_precision, dtype=self.w_map.dtype)

    def _g_matvec(self, v: Array) -> Array:
        """$G v$, jitter included."""
        if self.structure == "diag":
            return (self.ggn + self.jitter) * v
        return einx.dot("p q, q -> p", self.ggn, v) + self.jitter * v

    def precision(self, prior_precision: Array | float | None = None) -> Any:
        r"""The posterior precision $H_\lambda = G + \lambda I$.

        Args:
            prior_precision: $\lambda$; defaults to the posterior's own.

        Returns:
            A PSD-tagged ``lineax`` operator of size ``(P, P)``: a
            ``MatrixLinearOperator`` (full) or a ``DiagonalLinearOperator``
            (diag).
        """
        gaussx, lx = _require_gaussx()
        shift = self._lam(prior_precision) + self.jitter
        if self.structure == "diag":
            return lx.DiagonalLinearOperator(self.ggn + shift)
        p = self.w_map.shape[0]
        h = gaussx.symmetrize(self.ggn) + shift * jnp.eye(p, dtype=self.ggn.dtype)
        return lx.MatrixLinearOperator(h, lx.positive_semidefinite_tag)

    def mode(self, prior_precision: Array | float | None = None) -> Array:
        r"""$w_\lambda = H_\lambda^{-1} (G w^* + b)$, shape ``(P,)``.

        For a model linear in $w$ this is the exact posterior mean under
        $\lambda$, wherever $w^*$ was fitted.

        Args:
            prior_precision: $\lambda$; defaults to the posterior's own.

        Returns:
            The mode of the quadratic-model posterior.
        """
        gaussx, _ = _require_gaussx()
        rhs = self._g_matvec(self.w_map) + self.loglik_grad
        return gaussx.solve(self.precision(prior_precision), rhs)

    def predict(self, x: Array) -> tuple[Array, Array]:
        r"""Linearized predictive mean and variance of $f$ at ``x``.

        $$
        \mathbb{E}[f(x_*)] = f(x_*; w^*) + J_* (\bar w - w^*), \qquad
        \mathrm{Var}[f(x_*)] = J_* H_\lambda^{-1} J_*^\top,
        $$

        with $J_* = \partial f(x_*) / \partial w$ at $w^*$ and $\bar w$
        the posterior mean ``w_mean``. The variance is of the latent $f$;
        add ``sigma**2`` for the predictive of $y$.

        Args:
            x: Inputs, in whatever form ``predict_fn`` takes.

        Returns:
            ``(mean, var)``, each shaped like ``predict_fn(params, x)``.
        """
        gaussx, _ = _require_gaussx()
        _, unravel, rest = _split(self.params, self.param_filter)
        f = _forward(self.predict_fn, unravel, rest, x)
        f_map = self.predict_fn(self.params, x)
        jac = jax.jacfwd(f)(self.w_map)  # (M, P)
        shift = einx.dot("m p, p -> m", jac, self.w_mean - self.w_map)
        jac_t = einx.id("m p -> p m", jac)
        h = self.precision()
        if self.structure == "diag":
            # H⁻¹ is diagonal: Var = Σ_p J_*² / h_p.
            var = einx.sum("[p] m", jac_t * gaussx.solve_columns(h, jac_t))
        else:
            # H = L Lᵀ, so J_* H⁻¹ J_*ᵀ has diagonal Σ_p (L⁻¹ J_*ᵀ)².
            v = gaussx.solve_columns(gaussx.cholesky(h), jac_t)
            var = einx.sum("[p] m", v**2)
        mean = f_map + jnp.reshape(shift, f_map.shape)
        return mean, jnp.reshape(var, f_map.shape)

    def log_marginal_likelihood(
        self, prior_precision: Array | float | None = None
    ) -> Float[Array, ""]:
        r"""Laplace evidence $\log p(y \mid \lambda)$ of the quadratic model.

        $$
        \log p(y \mid \lambda) \approx \ell^* + b^\top \Delta
            - \tfrac{1}{2} \Delta^\top G \Delta
            - \tfrac{\lambda}{2} \lVert w_\lambda \rVert^2
            + \tfrac{P}{2} \log \lambda
            - \tfrac{1}{2} \log\det H_\lambda,
        \qquad \Delta = w_\lambda - w^*.
        $$

        The mode $w_\lambda$ is re-solved for each $\lambda$, so for a
        model linear in $w$ this is the exact Bayesian linear-regression
        evidence at every $\lambda$ (the jitter counts as part of $G$), not
        only at the $\lambda$ that produced $w^*$.

        Args:
            prior_precision: $\lambda$ to evaluate at; defaults to the
                posterior's own.

        Returns:
            Scalar log evidence.
        """
        gaussx, _ = _require_gaussx()
        lam = self._lam(prior_precision)
        w_lam = self.mode(lam)
        delta = w_lam - self.w_map
        p = self.w_map.shape[0]
        return (
            self.log_likelihood
            + jnp.dot(self.loglik_grad, delta)
            - 0.5 * jnp.dot(delta, self._g_matvec(delta))
            - 0.5 * lam * jnp.sum(w_lam**2)
            + 0.5 * p * jnp.log(lam)
            - 0.5 * gaussx.logdet(self.precision(lam))
        )

    def optimize_prior_precision(
        self, *, num_grid: int = 4097, num_refine: int = 100
    ) -> LastLayerLaplace:
        r"""Return a copy with $\lambda$ maximizing the Laplace evidence.

        In the eigenbasis $G = Q \operatorname{diag}(g) Q^\top$ (the
        identity for ``"diag"``), with $c = Q^\top (G w^* + b)$,
        $d = Q^\top w^*$, $\beta = Q^\top b$ and
        $\omega = c / (g + \lambda)$ the rotated mode, the evidence of
        `log_marginal_likelihood` is the $O(P)$ scalar function

        $$
        \ell^* + \beta^\top (\omega - d)
            - \tfrac{1}{2} \sum_i g_i (\omega_i - d_i)^2
            - \tfrac{\lambda}{2} \lVert \omega \rVert^2
            + \tfrac{P}{2} \log \lambda
            - \tfrac{1}{2} \sum_i \log (g_i + \lambda).
        $$

        It is maximized over $\log \lambda$ on a grid spanning half the
        dtype's exponent range each way (so $\lambda$ and every term stay
        finite, even when $w^* = 0$), then refined by golden-section search
        in the grid cell either side of the best point. For a model linear
        in $w$ this maximizes the exact evidence.

        Args:
            num_grid: Grid points on $\log \lambda$.
            num_refine: Golden-section steps (each shrinks the bracket by
                0.618).

        Returns:
            A new `LastLayerLaplace` with the tuned ``prior_precision`` and
            ``w_mean`` $= w_\lambda$ at it.
        """
        gaussx, lx = _require_gaussx()
        dtype = self.w_map.dtype
        if self.structure == "diag":
            g = self.ggn
            d, beta = self.w_map, self.loglik_grad
        else:
            g, q = gaussx.eig(
                lx.MatrixLinearOperator(
                    gaussx.symmetrize(self.ggn), lx.positive_semidefinite_tag
                )
            )
            d = einx.dot("p k, p -> k", q, self.w_map)
            beta = einx.dot("p k, p -> k", q, self.loglik_grad)
        # Round-off can leave eigenvalues of the PSD G slightly negative.
        g = jnp.clip(g.astype(dtype), 0.0) + self.jitter
        c = g * d + beta
        p = self.w_map.shape[0]

        def evidence(log_lam: Array) -> Array:
            lam = jnp.exp(log_lam)
            omega = c / (g + lam)
            delta = omega - d
            return (
                jnp.dot(beta, delta)
                - 0.5 * jnp.sum(g * delta**2)
                - 0.5 * lam * jnp.sum(omega**2)
                + 0.5 * p * log_lam
                - 0.5 * jnp.sum(jnp.log(g + lam))
            )

        bound = 0.5 * jnp.log(jnp.finfo(dtype).max)
        grid = jnp.linspace(-bound, bound, num_grid, dtype=dtype)
        best = jnp.argmax(jax.vmap(evidence)(grid))
        step = grid[1] - grid[0]
        lo = jnp.maximum(grid[best] - step, -bound)
        hi = jnp.minimum(grid[best] + step, bound)
        ratio = (jnp.sqrt(5.0) - 1.0) / 2.0

        def golden(_: int, bracket: tuple[Array, Array]) -> tuple[Array, Array]:
            a, b = bracket
            x1 = b - ratio * (b - a)
            x2 = a + ratio * (b - a)
            left = evidence(x1) >= evidence(x2)
            return jnp.where(left, a, x1), jnp.where(left, x2, b)

        lo, hi = jax.lax.fori_loop(0, num_refine, golden, (lo, hi))
        lam = jnp.exp(0.5 * (lo + hi)).astype(dtype)
        return eqx.tree_at(
            lambda m: (m.prior_precision, m.w_mean), self, (lam, self.mode(lam))
        )

precision(prior_precision: Array | float | None = None) -> Any

The posterior precision \(H_\lambda = G + \lambda I\).

Parameters:

Name	Type	Description	Default
prior_precision	Array | float | None	\(\lambda\); defaults to the posterior's own.	None

Returns:

Type	Description
Any	A PSD-tagged lineax operator of size (P, P): a
Any	MatrixLinearOperator (full) or a DiagonalLinearOperator
Any	(diag).

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

def precision(self, prior_precision: Array | float | None = None) -> Any:
    r"""The posterior precision $H_\lambda = G + \lambda I$.

    Args:
        prior_precision: $\lambda$; defaults to the posterior's own.

    Returns:
        A PSD-tagged ``lineax`` operator of size ``(P, P)``: a
        ``MatrixLinearOperator`` (full) or a ``DiagonalLinearOperator``
        (diag).
    """
    gaussx, lx = _require_gaussx()
    shift = self._lam(prior_precision) + self.jitter
    if self.structure == "diag":
        return lx.DiagonalLinearOperator(self.ggn + shift)
    p = self.w_map.shape[0]
    h = gaussx.symmetrize(self.ggn) + shift * jnp.eye(p, dtype=self.ggn.dtype)
    return lx.MatrixLinearOperator(h, lx.positive_semidefinite_tag)

mode(prior_precision: Array | float | None = None) -> Array

\(w_\lambda = H_\lambda^{-1} (G w^* + b)\), shape (P,).

For a model linear in \(w\) this is the exact posterior mean under \(\lambda\), wherever \(w^*\) was fitted.

Parameters:

Name	Type	Description	Default
prior_precision	Array | float | None	\(\lambda\); defaults to the posterior's own.	None

Returns:

Type	Description
Array	The mode of the quadratic-model posterior.

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

def mode(self, prior_precision: Array | float | None = None) -> Array:
    r"""$w_\lambda = H_\lambda^{-1} (G w^* + b)$, shape ``(P,)``.

    For a model linear in $w$ this is the exact posterior mean under
    $\lambda$, wherever $w^*$ was fitted.

    Args:
        prior_precision: $\lambda$; defaults to the posterior's own.

    Returns:
        The mode of the quadratic-model posterior.
    """
    gaussx, _ = _require_gaussx()
    rhs = self._g_matvec(self.w_map) + self.loglik_grad
    return gaussx.solve(self.precision(prior_precision), rhs)

predict(x: Array) -> tuple[Array, Array]

Linearized predictive mean and variance of \(f\) at x.

\[ \mathbb{E}[f(x_*)] = f(x_*; w^*) + J_* (\bar w - w^*), \qquad \mathrm{Var}[f(x_*)] = J_* H_\lambda^{-1} J_*^\top, \]

with \(J_* = \partial f(x_*) / \partial w\) at \(w^*\) and \(\bar w\) the posterior mean w_mean. The variance is of the latent \(f\); add sigma**2 for the predictive of \(y\).

Parameters:

Name	Type	Description	Default
x	Array	Inputs, in whatever form predict_fn takes.	required

Returns:

Type	Description
tuple[Array, Array]	(mean, var), each shaped like predict_fn(params, x).

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

def predict(self, x: Array) -> tuple[Array, Array]:
    r"""Linearized predictive mean and variance of $f$ at ``x``.

    $$
    \mathbb{E}[f(x_*)] = f(x_*; w^*) + J_* (\bar w - w^*), \qquad
    \mathrm{Var}[f(x_*)] = J_* H_\lambda^{-1} J_*^\top,
    $$

    with $J_* = \partial f(x_*) / \partial w$ at $w^*$ and $\bar w$
    the posterior mean ``w_mean``. The variance is of the latent $f$;
    add ``sigma**2`` for the predictive of $y$.

    Args:
        x: Inputs, in whatever form ``predict_fn`` takes.

    Returns:
        ``(mean, var)``, each shaped like ``predict_fn(params, x)``.
    """
    gaussx, _ = _require_gaussx()
    _, unravel, rest = _split(self.params, self.param_filter)
    f = _forward(self.predict_fn, unravel, rest, x)
    f_map = self.predict_fn(self.params, x)
    jac = jax.jacfwd(f)(self.w_map)  # (M, P)
    shift = einx.dot("m p, p -> m", jac, self.w_mean - self.w_map)
    jac_t = einx.id("m p -> p m", jac)
    h = self.precision()
    if self.structure == "diag":
        # H⁻¹ is diagonal: Var = Σ_p J_*² / h_p.
        var = einx.sum("[p] m", jac_t * gaussx.solve_columns(h, jac_t))
    else:
        # H = L Lᵀ, so J_* H⁻¹ J_*ᵀ has diagonal Σ_p (L⁻¹ J_*ᵀ)².
        v = gaussx.solve_columns(gaussx.cholesky(h), jac_t)
        var = einx.sum("[p] m", v**2)
    mean = f_map + jnp.reshape(shift, f_map.shape)
    return mean, jnp.reshape(var, f_map.shape)

log_marginal_likelihood(prior_precision: Array | float | None = None) -> Float[Array, '']

Laplace evidence \(\log p(y \mid \lambda)\) of the quadratic model.

\[ \log p(y \mid \lambda) \approx \ell^* + b^\top \Delta - \tfrac{1}{2} \Delta^\top G \Delta - \tfrac{\lambda}{2} \lVert w_\lambda \rVert^2 + \tfrac{P}{2} \log \lambda - \tfrac{1}{2} \log\det H_\lambda, \qquad \Delta = w_\lambda - w^*. \]

The mode \(w_\lambda\) is re-solved for each \(\lambda\), so for a model linear in \(w\) this is the exact Bayesian linear-regression evidence at every \(\lambda\) (the jitter counts as part of \(G\)), not only at the \(\lambda\) that produced \(w^*\).

Parameters:

Name	Type	Description	Default
prior_precision	Array | float | None	\(\lambda\) to evaluate at; defaults to the posterior's own.	None

Returns:

Type	Description
Float[Array, '']	Scalar log evidence.

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

def log_marginal_likelihood(
    self, prior_precision: Array | float | None = None
) -> Float[Array, ""]:
    r"""Laplace evidence $\log p(y \mid \lambda)$ of the quadratic model.

    $$
    \log p(y \mid \lambda) \approx \ell^* + b^\top \Delta
        - \tfrac{1}{2} \Delta^\top G \Delta
        - \tfrac{\lambda}{2} \lVert w_\lambda \rVert^2
        + \tfrac{P}{2} \log \lambda
        - \tfrac{1}{2} \log\det H_\lambda,
    \qquad \Delta = w_\lambda - w^*.
    $$

    The mode $w_\lambda$ is re-solved for each $\lambda$, so for a
    model linear in $w$ this is the exact Bayesian linear-regression
    evidence at every $\lambda$ (the jitter counts as part of $G$), not
    only at the $\lambda$ that produced $w^*$.

    Args:
        prior_precision: $\lambda$ to evaluate at; defaults to the
            posterior's own.

    Returns:
        Scalar log evidence.
    """
    gaussx, _ = _require_gaussx()
    lam = self._lam(prior_precision)
    w_lam = self.mode(lam)
    delta = w_lam - self.w_map
    p = self.w_map.shape[0]
    return (
        self.log_likelihood
        + jnp.dot(self.loglik_grad, delta)
        - 0.5 * jnp.dot(delta, self._g_matvec(delta))
        - 0.5 * lam * jnp.sum(w_lam**2)
        + 0.5 * p * jnp.log(lam)
        - 0.5 * gaussx.logdet(self.precision(lam))
    )

optimize_prior_precision(*, num_grid: int = 4097, num_refine: int = 100) -> LastLayerLaplace

Return a copy with \(\lambda\) maximizing the Laplace evidence.

In the eigenbasis \(G = Q \operatorname{diag}(g) Q^\top\) (the identity for "diag"), with \(c = Q^\top (G w^* + b)\), \(d = Q^\top w^*\), \(\beta = Q^\top b\) and \(\omega = c / (g + \lambda)\) the rotated mode, the evidence of log_marginal_likelihood is the \(O(P)\) scalar function

\[ \ell^* + \beta^\top (\omega - d) - \tfrac{1}{2} \sum_i g_i (\omega_i - d_i)^2 - \tfrac{\lambda}{2} \lVert \omega \rVert^2 + \tfrac{P}{2} \log \lambda - \tfrac{1}{2} \sum_i \log (g_i + \lambda). \]

It is maximized over \(\log \lambda\) on a grid spanning half the dtype's exponent range each way (so \(\lambda\) and every term stay finite, even when \(w^* = 0\)), then refined by golden-section search in the grid cell either side of the best point. For a model linear in \(w\) this maximizes the exact evidence.

Parameters:

Name	Type	Description	Default
num_grid	int	Grid points on \(\log \lambda\).	4097
num_refine	int	Golden-section steps (each shrinks the bracket by 0.618).	100

Returns:

Type	Description
LastLayerLaplace	A new LastLayerLaplace with the tuned prior_precision and
LastLayerLaplace	w_mean \(= w_\lambda\) at it.

Source code in packages/pyrox/src/pyrox/inference/_laplace.py

def optimize_prior_precision(
    self, *, num_grid: int = 4097, num_refine: int = 100
) -> LastLayerLaplace:
    r"""Return a copy with $\lambda$ maximizing the Laplace evidence.

    In the eigenbasis $G = Q \operatorname{diag}(g) Q^\top$ (the
    identity for ``"diag"``), with $c = Q^\top (G w^* + b)$,
    $d = Q^\top w^*$, $\beta = Q^\top b$ and
    $\omega = c / (g + \lambda)$ the rotated mode, the evidence of
    `log_marginal_likelihood` is the $O(P)$ scalar function

    $$
    \ell^* + \beta^\top (\omega - d)
        - \tfrac{1}{2} \sum_i g_i (\omega_i - d_i)^2
        - \tfrac{\lambda}{2} \lVert \omega \rVert^2
        + \tfrac{P}{2} \log \lambda
        - \tfrac{1}{2} \sum_i \log (g_i + \lambda).
    $$

    It is maximized over $\log \lambda$ on a grid spanning half the
    dtype's exponent range each way (so $\lambda$ and every term stay
    finite, even when $w^* = 0$), then refined by golden-section search
    in the grid cell either side of the best point. For a model linear
    in $w$ this maximizes the exact evidence.

    Args:
        num_grid: Grid points on $\log \lambda$.
        num_refine: Golden-section steps (each shrinks the bracket by
            0.618).

    Returns:
        A new `LastLayerLaplace` with the tuned ``prior_precision`` and
        ``w_mean`` $= w_\lambda$ at it.
    """
    gaussx, lx = _require_gaussx()
    dtype = self.w_map.dtype
    if self.structure == "diag":
        g = self.ggn
        d, beta = self.w_map, self.loglik_grad
    else:
        g, q = gaussx.eig(
            lx.MatrixLinearOperator(
                gaussx.symmetrize(self.ggn), lx.positive_semidefinite_tag
            )
        )
        d = einx.dot("p k, p -> k", q, self.w_map)
        beta = einx.dot("p k, p -> k", q, self.loglik_grad)
    # Round-off can leave eigenvalues of the PSD G slightly negative.
    g = jnp.clip(g.astype(dtype), 0.0) + self.jitter
    c = g * d + beta
    p = self.w_map.shape[0]

    def evidence(log_lam: Array) -> Array:
        lam = jnp.exp(log_lam)
        omega = c / (g + lam)
        delta = omega - d
        return (
            jnp.dot(beta, delta)
            - 0.5 * jnp.sum(g * delta**2)
            - 0.5 * lam * jnp.sum(omega**2)
            + 0.5 * p * log_lam
            - 0.5 * jnp.sum(jnp.log(g + lam))
        )

    bound = 0.5 * jnp.log(jnp.finfo(dtype).max)
    grid = jnp.linspace(-bound, bound, num_grid, dtype=dtype)
    best = jnp.argmax(jax.vmap(evidence)(grid))
    step = grid[1] - grid[0]
    lo = jnp.maximum(grid[best] - step, -bound)
    hi = jnp.minimum(grid[best] + step, bound)
    ratio = (jnp.sqrt(5.0) - 1.0) / 2.0

    def golden(_: int, bracket: tuple[Array, Array]) -> tuple[Array, Array]:
        a, b = bracket
        x1 = b - ratio * (b - a)
        x2 = a + ratio * (b - a)
        left = evidence(x1) >= evidence(x2)
        return jnp.where(left, a, x1), jnp.where(left, x2, b)

    lo, hi = jax.lax.fori_loop(0, num_refine, golden, (lo, hi))
    lam = jnp.exp(0.5 * (lo + hi)).astype(dtype)
    return eqx.tree_at(
        lambda m: (m.prior_precision, m.w_mean), self, (lam, self.mode(lam))
    )

Optimizer helpers

param_group_optimizer(groups: dict[str, optax.GradientTransformation], label_fn: Callable[[tuple, Any], str]) -> optax.GradientTransformation

Apply a different optimizer to each labelled parameter group.

Models that mix a large block of free latent parameters with a few kernel hyperparameters usually want different step sizes for each: the latents sit on a well-conditioned objective and tolerate a large step, while lengthscales and noise live on a log scale and destabilize under one. A 10x ratio is a common starting point.

A label_fn returning a label that is not a key of groups fails loudly at init time — optax.multi_transform raises a ValueError naming the offending labels.

Label by path, never by leaf value

optax.multi_transform evaluates a callable param_labels on the parameter tree at init and on the update tree at every update. A label_fn that inspects leaf values (a sign test, say) can therefore assign one group at init and another afterwards, silently applying the wrong transform or tripping a masked-state structure error. Depend only on path and on update-invariant leaf metadata such as shape / dtype, which are identical for parameters and their updates.

Parameters:

Name	Type	Description	Default
groups	dict[str, GradientTransformation]	Maps a group label to the optimizer for that group. Every label returned by label_fn must be a key here.	required
label_fn	Callable[[tuple, Any], str]	Called as label_fn(path, leaf) for each parameter, where path is the jax.tree_util key path. Returns the group label for that parameter. Must be a function of path (or update-invariant leaf metadata) only — see the warning above.	required

Returns:

Type	Description
GradientTransformation	A single optax.GradientTransformation suitable for
GradientTransformation	pyrox.inference.EnsembleMAP or numpyro.infer.SVI.

Examples:

>>> import optax
>>> def by_name(path, _):
...     return "latents" if "Z_T" in str(path) else "globals"
>>> opt = param_group_optimizer(
...     {"latents": optax.adam(1e-2), "globals": optax.adam(1e-3)},
...     by_name,
... )

Source code in packages/pyrox/src/pyrox/inference/_param_groups.py

def param_group_optimizer(
    groups: dict[str, optax.GradientTransformation],
    label_fn: Callable[[tuple, Any], str],
) -> optax.GradientTransformation:
    """Apply a different optimizer to each labelled parameter group.

    Models that mix a large block of free latent parameters with a few kernel
    hyperparameters usually want different step sizes for each: the latents
    sit on a well-conditioned objective and tolerate a large step, while
    lengthscales and noise live on a log scale and destabilize under one. A
    10x ratio is a common starting point.

    A ``label_fn`` returning a label that is not a key of ``groups`` fails
    loudly at ``init`` time — ``optax.multi_transform`` raises a
    ``ValueError`` naming the offending labels.

    !!! warning "Label by path, never by leaf value"
        ``optax.multi_transform`` evaluates a callable ``param_labels`` on
        the **parameter** tree at ``init`` and on the **update** tree at
        every ``update``. A ``label_fn`` that inspects leaf *values* (a
        sign test, say) can therefore assign one group at init and another
        afterwards, silently applying the wrong transform or tripping a
        masked-state structure error. Depend only on ``path`` and on
        update-invariant leaf metadata such as ``shape`` / ``dtype``, which
        are identical for parameters and their updates.

    Args:
        groups: Maps a group label to the optimizer for that group. Every
            label returned by ``label_fn`` must be a key here.
        label_fn: Called as ``label_fn(path, leaf)`` for each parameter,
            where ``path`` is the ``jax.tree_util`` key path. Returns the
            group label for that parameter. Must be a function of ``path``
            (or update-invariant leaf metadata) only — see the warning
            above.

    Returns:
        A single ``optax.GradientTransformation`` suitable for
        `pyrox.inference.EnsembleMAP` or ``numpyro.infer.SVI``.

    Examples:
        >>> import optax
        >>> def by_name(path, _):
        ...     return "latents" if "Z_T" in str(path) else "globals"
        >>> opt = param_group_optimizer(
        ...     {"latents": optax.adam(1e-2), "globals": optax.adam(1e-3)},
        ...     by_name,
        ... )
    """
    optax_mod = _require_optax()
    return optax_mod.multi_transform(
        groups,
        param_labels=lambda tree: jax.tree_util.tree_map_with_path(label_fn, tree),
    )

Result containers

EnsembleState

Bases: NamedTuple

Stacked state for an ensemble of optimizer runs.

Attributes:

Name	Type	Description
params	PyTree	Stacked parameter PyTree. Array leaves carry a leading (E,) axis; non-array leaves (e.g. captured activation functions inside an equinox.Module) are shared.
opt_state	PyTree	Stacked optax optimizer state with the same axis convention.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

class EnsembleState(NamedTuple):
    """Stacked state for an ensemble of optimizer runs.

    Attributes:
        params: Stacked parameter PyTree. Array leaves carry a leading
            ``(E,)`` axis; non-array leaves (e.g. captured activation
            functions inside an `equinox.Module`) are shared.
        opt_state: Stacked optax optimizer state with the same axis
            convention.
    """

    params: PyTree
    opt_state: PyTree

EnsembleResult

Bases: NamedTuple

Output of a full EnsembleMAP.run / EnsembleVI.run.

Attributes:

Name	Type	Description
params	PyTree	Final stacked parameters with leading (E,) axis.
losses	Float[Array, 'E T']	(E, num_epochs) per-step loss history.

Source code in packages/pyrox/src/pyrox/inference/_ensemble.py

class EnsembleResult(NamedTuple):
    """Output of a full `EnsembleMAP.run` / `EnsembleVI.run`.

    Attributes:
        params: Final stacked parameters with leading ``(E,)`` axis.
        losses: ``(E, num_epochs)`` per-step loss history.
    """

    params: PyTree
    losses: Float[Array, "E T"]