Quadrature & Moment Matching

Gaussian integration: \(\mathbb{E}_{x \sim \mathcal{N}(\mu, \Sigma)}[f(x)]\) via deterministic rules (Gauss-Hermite, unscented / cubature, Taylor) or Monte Carlo, behind one AbstractIntegrator interface. Everything that needs an expectation — expected log-likelihoods, EP tilted moments, uncertain-input GP predictions — takes an integrator argument, so swapping the rule never touches the model.

State & integrators

GaussianState pairs a mean with a covariance operator; integrators propagate functions of it.

Structured linear algebra and Gaussian primitives for JAX.

GaussianState

Bases: Module

Gaussian distribution as (mean, covariance operator) pair.

Attributes:

Name	Type	Description
mean	Float[Array, ' N']	Mean vector, shape (N,).
cov	AbstractLinearOperator	Covariance operator, shape (N, N).

Source code in src/gaussx/_quadrature/_types.py

class GaussianState(eqx.Module):
    """Gaussian distribution as (mean, covariance operator) pair.

    Attributes:
        mean: Mean vector, shape ``(N,)``.
        cov: Covariance operator, shape ``(N, N)``.
    """

    mean: Float[Array, " N"]
    cov: lx.AbstractLinearOperator

PropagationResult

Bases: Module

Output of uncertainty propagation through a nonlinear function.

Tag convention for state.cov: lineax.positive_semidefinite_tag when the rule guarantees a PSD covariance at the input dimension (integrator.guarantees_psd(N): Gauss-Hermite, spherical cubature, unscented with non-negative weights, Monte Carlo, ADF, Taylor), else lineax.symmetric_tag (negative-weight rules such as the default scaled unscented transform or the fifth-order cubature rule). Dispatch on lineax.is_positive_semidefinite therefore depends only on whether PSD is guaranteed, not on which rule produced the state.

Attributes:

Name	Type	Description
state	GaussianState	Output Gaussian distribution.
cross_cov	Float[Array, 'N_in N_out'] | None	Input-output cross-covariance, shape (N_in, N_out). Used for downstream Kalman updates. None if not computed.

Source code in src/gaussx/_quadrature/_types.py

class PropagationResult(eqx.Module):
    """Output of uncertainty propagation through a nonlinear function.

    Tag convention for ``state.cov``: `lineax.positive_semidefinite_tag`
    when the rule guarantees a PSD covariance at the input dimension
    (``integrator.guarantees_psd(N)``: Gauss-Hermite, spherical cubature,
    unscented with non-negative weights, Monte Carlo, ADF, Taylor), else
    `lineax.symmetric_tag` (negative-weight rules such as the default
    scaled unscented transform or the fifth-order cubature rule).
    Dispatch on `lineax.is_positive_semidefinite` therefore depends only
    on whether PSD is guaranteed, not on which rule produced the state.

    Attributes:
        state: Output Gaussian distribution.
        cross_cov: Input-output cross-covariance, shape ``(N_in, N_out)``.
            Used for downstream Kalman updates. ``None`` if not computed.
    """

    state: GaussianState
    cross_cov: Float[Array, "N_in N_out"] | None

AbstractIntegrator

Bases: Module

Protocol for Gaussian integral approximation.

Subclasses implement integrate to propagate a Gaussian through a nonlinear function, returning an approximate output distribution and (optionally) input-output cross-covariance.

Source code in src/gaussx/_quadrature/_integrator.py

class AbstractIntegrator(eqx.Module):
    """Protocol for Gaussian integral approximation.

    Subclasses implement ``integrate`` to propagate a Gaussian through
    a nonlinear function, returning an approximate output distribution
    and (optionally) input-output cross-covariance.
    """

    @abc.abstractmethod
    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate a Gaussian through ``fn``, returning output moments.

        Args:
            fn: Nonlinear function mapping ``(N,) -> (M,)``.
            state: Input Gaussian distribution.

        Returns:
            ``PropagationResult`` with output distribution and optional
            cross-covariance.
        """
        ...

    def guarantees_psd(self, dim: int) -> bool:
        """Whether the rule's moment triple is PSD by construction at ``dim``.

        A rule whose covariance weights are all non-negative (or a
        linearisation, which propagates the covariance exactly) returns a
        consistent joint over ``(x, g(x))``, so its output covariance and
        the implied conditional are PSD. Consumers such as
        `gaussx.nonlinear_kalman_filter` skip their definiteness checks for
        such rules (gh-331). The default is ``False``, so a custom rule is
        always checked; the built-in positive-weight rules override it.

        Args:
            dim: Input dimension ``N`` the rule is applied at.

        Returns:
            ``True`` if every covariance weight is non-negative at ``dim``.
        """
        del dim
        return False

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """Return the raw quadrature points and weights for ``state``.

        Point-based rules override this so that consumers needing the raw
        evaluations — `gaussx.moment_match`,
        `gaussx.statistical_linear_regression` — can reuse a single pass
        over the points instead of re-deriving the rule.

        Args:
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(points, w_m, w_c)`` where ``points`` has shape
            ``(P, N)`` in the input space, ``w_m`` are the mean weights and
            ``w_c`` the covariance weights, both shape ``(P,)``.

        Raises:
            NotImplementedError: If the integrator is not point-based, e.g.
                `TaylorIntegrator`, which linearises instead of sampling.
        """
        msg = (
            f"{type(self).__name__} is not a point-based rule and does not "
            f"expose quadrature points and weights."
        )
        raise NotImplementedError(msg)

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult abstractmethod

Propagate a Gaussian through fn, returning output moments.

Parameters:

Name	Type	Description	Default
fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	Nonlinear function mapping (N,) -> (M,).	required
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
PropagationResult	PropagationResult with output distribution and optional
PropagationResult	cross-covariance.

Source code in src/gaussx/_quadrature/_integrator.py

@abc.abstractmethod
def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate a Gaussian through ``fn``, returning output moments.

    Args:
        fn: Nonlinear function mapping ``(N,) -> (M,)``.
        state: Input Gaussian distribution.

    Returns:
        ``PropagationResult`` with output distribution and optional
        cross-covariance.
    """
    ...

guarantees_psd(dim: int) -> bool

Whether the rule's moment triple is PSD by construction at dim.

A rule whose covariance weights are all non-negative (or a linearisation, which propagates the covariance exactly) returns a consistent joint over (x, g(x)), so its output covariance and the implied conditional are PSD. Consumers such as gaussx.nonlinear_kalman_filter skip their definiteness checks for such rules (gh-331). The default is False, so a custom rule is always checked; the built-in positive-weight rules override it.

Parameters:

Name	Type	Description	Default
dim	int	Input dimension N the rule is applied at.	required

Returns:

Type	Description
bool	True if every covariance weight is non-negative at dim.

Source code in src/gaussx/_quadrature/_integrator.py

def guarantees_psd(self, dim: int) -> bool:
    """Whether the rule's moment triple is PSD by construction at ``dim``.

    A rule whose covariance weights are all non-negative (or a
    linearisation, which propagates the covariance exactly) returns a
    consistent joint over ``(x, g(x))``, so its output covariance and
    the implied conditional are PSD. Consumers such as
    `gaussx.nonlinear_kalman_filter` skip their definiteness checks for
    such rules (gh-331). The default is ``False``, so a custom rule is
    always checked; the built-in positive-weight rules override it.

    Args:
        dim: Input dimension ``N`` the rule is applied at.

    Returns:
        ``True`` if every covariance weight is non-negative at ``dim``.
    """
    del dim
    return False

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Return the raw quadrature points and weights for state.

Point-based rules override this so that consumers needing the raw evaluations — gaussx.moment_match, gaussx.statistical_linear_regression — can reuse a single pass over the points instead of re-deriving the rule.

Parameters:

Name	Type	Description	Default
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (points, w_m, w_c) where points has shape
Float[Array, ' P']	(P, N) in the input space, w_m are the mean weights and
Float[Array, ' P']	w_c the covariance weights, both shape (P,).

Raises:

Type	Description
NotImplementedError	If the integrator is not point-based, e.g. TaylorIntegrator, which linearises instead of sampling.

Source code in src/gaussx/_quadrature/_integrator.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """Return the raw quadrature points and weights for ``state``.

    Point-based rules override this so that consumers needing the raw
    evaluations — `gaussx.moment_match`,
    `gaussx.statistical_linear_regression` — can reuse a single pass
    over the points instead of re-deriving the rule.

    Args:
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(points, w_m, w_c)`` where ``points`` has shape
        ``(P, N)`` in the input space, ``w_m`` are the mean weights and
        ``w_c`` the covariance weights, both shape ``(P,)``.

    Raises:
        NotImplementedError: If the integrator is not point-based, e.g.
            `TaylorIntegrator`, which linearises instead of sampling.
    """
    msg = (
        f"{type(self).__name__} is not a point-based rule and does not "
        f"expose quadrature points and weights."
    )
    raise NotImplementedError(msg)

GaussHermiteIntegrator

Bases: AbstractIntegrator

Gauss-Hermite quadrature integrator.

Approximates Gaussian expectations using tensor-product Gauss-Hermite quadrature:

E[g(f)] \approx \sum_i w_i \cdot g(\mu + L z_i)

where (z_i, w_i) are GH points/weights in standard normal space and L is the square root of the covariance.

Exact for polynomials up to degree 2 * order - 1.

Cost: a tensor-product rule with P = order^dim points, each one an evaluation of the integrand. At the default order=20 that is 20 points at dim=1, 400 at dim=2, 8 000 at dim=3 and 160 000 at dim=4. Use it for dim <= 2 (or lower order); for larger dim prefer CubatureIntegrator (2 dim points) or FifthOrderCubatureIntegrator (2 dim^2 + 1 points).

Attributes:

Name	Type	Description
order	int	Number of quadrature points per dimension. Default 20.

Source code in src/gaussx/_quadrature/_gauss_hermite.py

class GaussHermiteIntegrator(AbstractIntegrator):
    r"""Gauss-Hermite quadrature integrator.

    Approximates Gaussian expectations using tensor-product Gauss-Hermite
    quadrature:

        E[g(f)] \approx \sum_i w_i \cdot g(\mu + L z_i)

    where ``(z_i, w_i)`` are GH points/weights in standard normal space
    and ``L`` is the square root of the covariance.

    Exact for polynomials up to degree ``2 * order - 1``.

    Cost: a tensor-product rule with ``P = order^dim`` points, each one an
    evaluation of the integrand. At the default ``order=20`` that is 20
    points at ``dim=1``, 400 at ``dim=2``, 8 000 at ``dim=3`` and 160 000 at
    ``dim=4``. Use it for ``dim <= 2`` (or lower ``order``); for larger
    ``dim`` prefer `CubatureIntegrator` (``2 dim`` points) or
    `FifthOrderCubatureIntegrator` (``2 dim^2 + 1`` points).

    Attributes:
        order: Number of quadrature points per dimension. Default ``20``.
    """

    order: int = eqx.field(static=True, default=20)

    def guarantees_psd(self, dim: int) -> bool:
        """Gauss-Hermite weights are positive."""
        del dim
        return True

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via Gauss-Hermite quadrature."""
        chi, w, _ = self.points_and_weights(state)

        # Propagate all quadrature points
        Y = jax.vmap(fn)(chi)  # (P, M)

        return assemble_propagation_result(
            chi, Y, state.mean, w, psd=self.guarantees_psd(state.mean.shape[0])
        )

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """Return the tensor-product Gauss-Hermite points and weights.

        Args:
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(points, w_m, w_c)`` with ``P = order^N``. The rule
            uses a single weight set, so ``w_m`` and ``w_c`` are the same
            array, normalised to sum to one.
        """

        mu = state.mean
        N = mu.shape[0]

        # GH points in standard normal space, in the input dtype so a
        # float32 belief is not promoted under x64.
        z, w = gauss_hermite_points(self.order, N, dtype=mu.dtype)

        # Transform to input space: xᵢ = μ + S zᵢ
        S = sqrt(state.cov).as_matrix()
        chi = mu[None, :] + z @ S.T  # (P, N)

        # Normalize weights to sum to 1
        w = w / jnp.sum(w)

        return chi, w, w

guarantees_psd(dim: int) -> bool

Gauss-Hermite weights are positive.

Source code in src/gaussx/_quadrature/_gauss_hermite.py

def guarantees_psd(self, dim: int) -> bool:
    """Gauss-Hermite weights are positive."""
    del dim
    return True

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via Gauss-Hermite quadrature.

Source code in src/gaussx/_quadrature/_gauss_hermite.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via Gauss-Hermite quadrature."""
    chi, w, _ = self.points_and_weights(state)

    # Propagate all quadrature points
    Y = jax.vmap(fn)(chi)  # (P, M)

    return assemble_propagation_result(
        chi, Y, state.mean, w, psd=self.guarantees_psd(state.mean.shape[0])
    )

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Return the tensor-product Gauss-Hermite points and weights.

Parameters:

Name	Type	Description	Default
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (points, w_m, w_c) with P = order^N. The rule
Float[Array, ' P']	uses a single weight set, so w_m and w_c are the same
Float[Array, ' P']	array, normalised to sum to one.

Source code in src/gaussx/_quadrature/_gauss_hermite.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """Return the tensor-product Gauss-Hermite points and weights.

    Args:
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(points, w_m, w_c)`` with ``P = order^N``. The rule
        uses a single weight set, so ``w_m`` and ``w_c`` are the same
        array, normalised to sum to one.
    """

    mu = state.mean
    N = mu.shape[0]

    # GH points in standard normal space, in the input dtype so a
    # float32 belief is not promoted under x64.
    z, w = gauss_hermite_points(self.order, N, dtype=mu.dtype)

    # Transform to input space: xᵢ = μ + S zᵢ
    S = sqrt(state.cov).as_matrix()
    chi = mu[None, :] + z @ S.T  # (P, N)

    # Normalize weights to sum to 1
    w = w / jnp.sum(w)

    return chi, w, w

UnscentedIntegrator

Bases: AbstractIntegrator

Unscented transform: deterministic sigma points.

Generates 2N+1 sigma points around the mean, propagates them through the nonlinear function, and reconstructs output moments:

chi_i = mu + sqrt((N + lambda) * Sigma) @ xi_i
y_i = f(chi_i)
mu_y = sum(w_m * y_i)
Sigma_y = sum(w_c * (y_i - mu_y)(y_i - mu_y)^T)
cross_cov = sum(w_c * (chi_i - mu)(y_i - mu_y)^T)

where lambda = alpha^2 * (N + kappa) - N.

Warning

The default alpha=1e-3 (the classic Wan-van der Merwe scaled transform) puts a centre mean weight of about -1e6 on the points and recovers the moments by cancellation, losing about six digits: in float32 a smooth nonlinearity is wrong in the second digit (gh-310). Prefer UnscentedIntegrator(alpha=1.0), the symmetric 2N+1 rule with a zero centre mean weight (exact for affine maps and positive covariance weights), which is what moment_transform and the nonlinear Kalman filters use by default. A UserWarning is emitted when the centre mean weight reaches magnitude 1e3 with a float32 state.

Attributes:

Name	Type	Description
alpha	float	Spread parameter. Default 1e-3; 1.0 recommended.
beta	float	Prior knowledge parameter (2.0 optimal for Gaussian).
kappa	float	Secondary scaling. Default 0.0.

Source code in src/gaussx/_quadrature/_unscented.py

class UnscentedIntegrator(AbstractIntegrator):
    r"""Unscented transform: deterministic sigma points.

    Generates ``2N+1`` sigma points around the mean, propagates them
    through the nonlinear function, and reconstructs output moments:

        chi_i = mu + sqrt((N + lambda) * Sigma) @ xi_i
        y_i = f(chi_i)
        mu_y = sum(w_m * y_i)
        Sigma_y = sum(w_c * (y_i - mu_y)(y_i - mu_y)^T)
        cross_cov = sum(w_c * (chi_i - mu)(y_i - mu_y)^T)

    where ``lambda = alpha^2 * (N + kappa) - N``.

    Warning:
        The default ``alpha=1e-3`` (the classic Wan-van der Merwe scaled
        transform) puts a centre mean weight of about ``-1e6`` on the
        points and recovers the moments by cancellation, losing about six
        digits: in float32 a smooth nonlinearity is wrong in the second
        digit (gh-310). Prefer ``UnscentedIntegrator(alpha=1.0)``, the
        symmetric ``2N+1`` rule with a zero centre mean weight (exact for
        affine maps and positive covariance weights), which is what
        `moment_transform` and the nonlinear Kalman filters use by
        default. A `UserWarning` is emitted when the centre
        mean weight reaches magnitude ``1e3`` with a float32 state.

    Attributes:
        alpha: Spread parameter. Default ``1e-3``; ``1.0`` recommended.
        beta: Prior knowledge parameter (2.0 optimal for Gaussian).
        kappa: Secondary scaling. Default ``0.0``.
    """

    alpha: float = eqx.field(static=True, default=1e-3)
    beta: float = eqx.field(static=True, default=2.0)
    kappa: float = eqx.field(static=True, default=0.0)

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via unscented transform."""
        chi, w_m, w_c = self.points_and_weights(state)
        Y = jax.vmap(fn)(chi)
        return assemble_propagation_result(
            chi, Y, state.mean, w_m, w_c, psd=self.guarantees_psd(state.mean.shape[0])
        )

    def guarantees_psd(self, dim: int) -> bool:
        """Whether every scaled-unscented covariance weight is non-negative.

        With ``lambda = alpha^2 (N + kappa) - N`` the weights are
        ``1 / (2 (N + lambda))`` off-centre and
        ``lambda / (N + lambda) + 1 - alpha^2 + beta`` at the centre, so
        small ``alpha`` (the ``1e-3`` default) makes the centre weight
        large and negative; ``alpha = 1`` does not.
        """
        lam = self.alpha**2 * (dim + self.kappa) - dim
        spread = dim + lam
        if spread <= 0:
            return False
        return lam / spread + 1.0 - self.alpha**2 + self.beta >= 0

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """Return the scaled unscented sigma points and weights.

        Args:
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(points, w_m, w_c)`` with ``P = 2N + 1``. Mean and
            covariance weights differ in the centre point.
        """
        return sigma_points(
            state.mean,
            state.cov,
            alpha=self.alpha,
            beta=self.beta,
            kappa=self.kappa,
        )

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via unscented transform.

Source code in src/gaussx/_quadrature/_unscented.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via unscented transform."""
    chi, w_m, w_c = self.points_and_weights(state)
    Y = jax.vmap(fn)(chi)
    return assemble_propagation_result(
        chi, Y, state.mean, w_m, w_c, psd=self.guarantees_psd(state.mean.shape[0])
    )

guarantees_psd(dim: int) -> bool

Whether every scaled-unscented covariance weight is non-negative.

With lambda = alpha^2 (N + kappa) - N the weights are 1 / (2 (N + lambda)) off-centre and lambda / (N + lambda) + 1 - alpha^2 + beta at the centre, so small alpha (the 1e-3 default) makes the centre weight large and negative; alpha = 1 does not.

Source code in src/gaussx/_quadrature/_unscented.py

def guarantees_psd(self, dim: int) -> bool:
    """Whether every scaled-unscented covariance weight is non-negative.

    With ``lambda = alpha^2 (N + kappa) - N`` the weights are
    ``1 / (2 (N + lambda))`` off-centre and
    ``lambda / (N + lambda) + 1 - alpha^2 + beta`` at the centre, so
    small ``alpha`` (the ``1e-3`` default) makes the centre weight
    large and negative; ``alpha = 1`` does not.
    """
    lam = self.alpha**2 * (dim + self.kappa) - dim
    spread = dim + lam
    if spread <= 0:
        return False
    return lam / spread + 1.0 - self.alpha**2 + self.beta >= 0

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Return the scaled unscented sigma points and weights.

Parameters:

Name	Type	Description	Default
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (points, w_m, w_c) with P = 2N + 1. Mean and
Float[Array, ' P']	covariance weights differ in the centre point.

Source code in src/gaussx/_quadrature/_unscented.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """Return the scaled unscented sigma points and weights.

    Args:
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(points, w_m, w_c)`` with ``P = 2N + 1``. Mean and
        covariance weights differ in the centre point.
    """
    return sigma_points(
        state.mean,
        state.cov,
        alpha=self.alpha,
        beta=self.beta,
        kappa=self.kappa,
    )

CubatureIntegrator

Bases: AbstractIntegrator

Third-order spherical-radial cubature.

The 2N-point, equally weighted rule of Arasaratnam & Haykin (2009), exact for polynomials up to total degree 3. Driving gaussx.nonlinear_kalman_filter with it gives the textbook cubature Kalman filter (CKF).

It is the degree-3 sibling of gaussx.FifthOrderCubatureIntegrator, and closely related to gaussx.UnscentedIntegrator: the unscented transform with alpha=1, beta=0, kappa=0 places its 2N non-centre points identically, differing only by the centre point and its weight. Unlike the scaled unscented transform, every weight here is positive (1 / 2N), so an assembled covariance cannot be dragged indefinite by a negative weight.

Source code in src/gaussx/_quadrature/_cubature.py

class CubatureIntegrator(AbstractIntegrator):
    r"""Third-order spherical-radial cubature.

    The ``2N``-point, equally weighted rule of Arasaratnam & Haykin
    (2009), exact for polynomials up to total degree 3. Driving
    `gaussx.nonlinear_kalman_filter` with it gives the textbook cubature
    Kalman filter (CKF).

    It is the degree-3 sibling of `gaussx.FifthOrderCubatureIntegrator`,
    and closely related to `gaussx.UnscentedIntegrator`: the unscented
    transform with ``alpha=1``, ``beta=0``, ``kappa=0`` places its ``2N``
    non-centre points identically, differing only by the centre point and
    its weight. Unlike the scaled unscented transform, every weight here is
    positive (``1 / 2N``), so an assembled covariance cannot be dragged
    indefinite by a negative weight.
    """

    def guarantees_psd(self, dim: int) -> bool:
        """All ``2N`` cubature weights are ``1 / 2N``."""
        del dim
        return True

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via third-order cubature."""
        chi, weights = cubature_points(state.mean, state.cov)
        Y = jax.vmap(fn)(chi)
        return assemble_propagation_result(
            chi, Y, state.mean, weights, psd=self.guarantees_psd(state.mean.shape[0])
        )

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """Return the third-order cubature points and weights.

        Args:
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(points, w_m, w_c)`` with ``P = 2N``. The rule uses a
            single weight set, so ``w_m`` and ``w_c`` are the same array.
        """
        chi, weights = cubature_points(state.mean, state.cov)
        return chi, weights, weights

guarantees_psd(dim: int) -> bool

All 2N cubature weights are 1 / 2N.

Source code in src/gaussx/_quadrature/_cubature.py

def guarantees_psd(self, dim: int) -> bool:
    """All ``2N`` cubature weights are ``1 / 2N``."""
    del dim
    return True

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via third-order cubature.

Source code in src/gaussx/_quadrature/_cubature.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via third-order cubature."""
    chi, weights = cubature_points(state.mean, state.cov)
    Y = jax.vmap(fn)(chi)
    return assemble_propagation_result(
        chi, Y, state.mean, weights, psd=self.guarantees_psd(state.mean.shape[0])
    )

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Return the third-order cubature points and weights.

Parameters:

Name	Type	Description	Default
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (points, w_m, w_c) with P = 2N. The rule uses a
Float[Array, ' P']	single weight set, so w_m and w_c are the same array.

Source code in src/gaussx/_quadrature/_cubature.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """Return the third-order cubature points and weights.

    Args:
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(points, w_m, w_c)`` with ``P = 2N``. The rule uses a
        single weight set, so ``w_m`` and ``w_c`` are the same array.
    """
    chi, weights = cubature_points(state.mean, state.cov)
    return chi, weights, weights

FifthOrderCubatureIntegrator

Bases: AbstractIntegrator

Fifth-order fully symmetric cubature (McNamee & Stenger).

Exact for polynomials up to total degree 5 in N variables, using 2 N^2 + 1 points. This is the practical middle ground between the degree-3 sigma-point rules and a tensor-product Gauss-Hermite rule:

================================== ============== ================ Rule Points Exact to degree ================================== ============== ================ UnscentedIntegrator 2N + 1 3 FifthOrderCubatureIntegrator 2 N^2 + 1 5 GaussHermiteIntegrator (order H) H^N 2H − 1 ================================== ============== ================

It is the default integrator in BayesNewton, and the lowest-degree rule that gets the curvature right for likelihoods whose log-density behaves like a quartic near the mode (Poisson, Bernoulli with heavy-tailed inputs).

Note

The rule carries negative weights on the axis shell for N > 4 — see gaussx.fifth_order_cubature_points. Moments through degree 5 remain exact, but an output covariance assembled from very few function evaluations is not guaranteed PSD in high dimension.

Source code in src/gaussx/_quadrature/_fifth_order.py

class FifthOrderCubatureIntegrator(AbstractIntegrator):
    r"""Fifth-order fully symmetric cubature (McNamee & Stenger).

    Exact for polynomials up to total degree 5 in ``N`` variables, using
    ``2 N^2 + 1`` points. This is the practical middle ground between the
    degree-3 sigma-point rules and a tensor-product Gauss-Hermite rule:

    ==================================  ==============  ================
    Rule                                Points          Exact to degree
    ==================================  ==============  ================
    `UnscentedIntegrator`               ``2N + 1``      3
    `FifthOrderCubatureIntegrator`      ``2 N^2 + 1``   5
    `GaussHermiteIntegrator` (order H)  ``H^N``         ``2H − 1``
    ==================================  ==============  ================

    It is the default integrator in BayesNewton, and the lowest-degree rule
    that gets the curvature right for likelihoods whose log-density behaves
    like a quartic near the mode (Poisson, Bernoulli with heavy-tailed
    inputs).

    Note:
        The rule carries negative weights on the axis shell for ``N > 4``
        — see `gaussx.fifth_order_cubature_points`. Moments through degree
        5 remain exact, but an output covariance assembled from very few
        function evaluations is not guaranteed PSD in high dimension.
    """

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via fifth-order cubature."""
        chi, weights = fifth_order_cubature_points(state.mean, state.cov)
        Y = jax.vmap(fn)(chi)
        return assemble_propagation_result(
            chi, Y, state.mean, weights, psd=self.guarantees_psd(state.mean.shape[0])
        )

    def guarantees_psd(self, dim: int) -> bool:
        """The axis-shell weights are negative for ``N > 4``."""
        return dim <= 4

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """Return the fifth-order cubature points and weights.

        Args:
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(points, w_m, w_c)``. The rule uses a single weight
            set, so ``w_m`` and ``w_c`` are the same array.
        """
        chi, weights = fifth_order_cubature_points(state.mean, state.cov)
        return chi, weights, weights

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via fifth-order cubature.

Source code in src/gaussx/_quadrature/_fifth_order.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via fifth-order cubature."""
    chi, weights = fifth_order_cubature_points(state.mean, state.cov)
    Y = jax.vmap(fn)(chi)
    return assemble_propagation_result(
        chi, Y, state.mean, weights, psd=self.guarantees_psd(state.mean.shape[0])
    )

guarantees_psd(dim: int) -> bool

The axis-shell weights are negative for N > 4.

Source code in src/gaussx/_quadrature/_fifth_order.py

def guarantees_psd(self, dim: int) -> bool:
    """The axis-shell weights are negative for ``N > 4``."""
    return dim <= 4

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Return the fifth-order cubature points and weights.

Parameters:

Name	Type	Description	Default
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (points, w_m, w_c). The rule uses a single weight
Float[Array, ' P']	set, so w_m and w_c are the same array.

Source code in src/gaussx/_quadrature/_fifth_order.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """Return the fifth-order cubature points and weights.

    Args:
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(points, w_m, w_c)``. The rule uses a single weight
        set, so ``w_m`` and ``w_c`` are the same array.
    """
    chi, weights = fifth_order_cubature_points(state.mean, state.cov)
    return chi, weights, weights

TaylorIntegrator

Bases: AbstractIntegrator

1st or 2nd order Taylor expansion for uncertainty propagation.

1st order (EKF):

mu_y = f(mu_x)
Sigma_y = J @ Sigma_x @ J^T
cross_cov = Sigma_x @ J^T

2nd order:

mu_y_i += 0.5 * tr(H_i @ Sigma_x)
Sigma_y += correction from Hessians

Attributes:

Name	Type	Description
order	int	Taylor expansion order (1 or 2). Default 1.
correct_variance	bool	If True and order=2, apply 2nd-order covariance correction using 4th Gaussian moments. Default True to preserve the historical order=2 behaviour. Set to False for the mean-only correction used in the standard EKF literature. Ignored when order=1.

Source code in src/gaussx/_quadrature/_taylor.py

class TaylorIntegrator(AbstractIntegrator):
    r"""1st or 2nd order Taylor expansion for uncertainty propagation.

    **1st order (EKF)**:

        mu_y = f(mu_x)
        Sigma_y = J @ Sigma_x @ J^T
        cross_cov = Sigma_x @ J^T

    **2nd order**:

        mu_y_i += 0.5 * tr(H_i @ Sigma_x)
        Sigma_y += correction from Hessians

    Attributes:
        order: Taylor expansion order (1 or 2). Default 1.
        correct_variance: If True and order=2, apply 2nd-order covariance
            correction using 4th Gaussian moments. Default True to preserve
            the historical ``order=2`` behaviour. Set to False for the
            mean-only correction used in the standard EKF literature.
            Ignored when order=1.
    """

    order: int = eqx.field(static=True, default=1)
    correct_variance: bool = eqx.field(static=True, default=True)

    def guarantees_psd(self, dim: int) -> bool:
        """A linearisation propagates a consistent joint (``J P Jᵀ`` plus a PSD
        second-order correction)."""
        del dim
        return True

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via Taylor expansion."""
        if self.order not in (1, 2):
            msg = f"TaylorIntegrator.order must be 1 or 2, got {self.order}"
            raise ValueError(msg)
        mu = state.mean
        Sigma = state.cov.as_matrix()

        # Evaluate function and Jacobian at the mean
        f_mu = fn(mu)
        J = jax.jacobian(fn)(mu)  # (M, N)

        if self.order == 1:
            mu_y = f_mu
            Sigma_y = J @ Sigma @ J.T
        else:
            # 2nd order correction: mu_y_i += 0.5 * tr(H_i @ Sigma_x)
            H_fn = jax.hessian(fn)
            H = H_fn(mu)  # (M, N, N)
            corrections = jax.vmap(lambda H_i: jnp.trace(H_i @ Sigma))(H)
            mu_y = f_mu + 0.5 * corrections
            Sigma_y = J @ Sigma @ J.T
            if self.correct_variance:
                second_order_cov = jax.vmap(
                    lambda H_i: jax.vmap(
                        lambda H_j: 0.5 * jnp.trace(H_i @ Sigma @ H_j @ Sigma)
                    )(H)
                )(H)
                Sigma_y = Sigma_y + second_order_cov

        # Symmetrize for numerical stability
        Sigma_y = symmetrize(Sigma_y)

        # Cross-covariance: Sigma_x @ J^T
        cross_cov = Sigma @ J.T

        cov_y = lx.MatrixLinearOperator(Sigma_y, lx.positive_semidefinite_tag)
        out_state = GaussianState(mean=mu_y, cov=cov_y)

        return PropagationResult(state=out_state, cross_cov=cross_cov)

guarantees_psd(dim: int) -> bool

A linearisation propagates a consistent joint (J P Jᵀ plus a PSD second-order correction).

Source code in src/gaussx/_quadrature/_taylor.py

def guarantees_psd(self, dim: int) -> bool:
    """A linearisation propagates a consistent joint (``J P Jᵀ`` plus a PSD
    second-order correction)."""
    del dim
    return True

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via Taylor expansion.

Source code in src/gaussx/_quadrature/_taylor.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via Taylor expansion."""
    if self.order not in (1, 2):
        msg = f"TaylorIntegrator.order must be 1 or 2, got {self.order}"
        raise ValueError(msg)
    mu = state.mean
    Sigma = state.cov.as_matrix()

    # Evaluate function and Jacobian at the mean
    f_mu = fn(mu)
    J = jax.jacobian(fn)(mu)  # (M, N)

    if self.order == 1:
        mu_y = f_mu
        Sigma_y = J @ Sigma @ J.T
    else:
        # 2nd order correction: mu_y_i += 0.5 * tr(H_i @ Sigma_x)
        H_fn = jax.hessian(fn)
        H = H_fn(mu)  # (M, N, N)
        corrections = jax.vmap(lambda H_i: jnp.trace(H_i @ Sigma))(H)
        mu_y = f_mu + 0.5 * corrections
        Sigma_y = J @ Sigma @ J.T
        if self.correct_variance:
            second_order_cov = jax.vmap(
                lambda H_i: jax.vmap(
                    lambda H_j: 0.5 * jnp.trace(H_i @ Sigma @ H_j @ Sigma)
                )(H)
            )(H)
            Sigma_y = Sigma_y + second_order_cov

    # Symmetrize for numerical stability
    Sigma_y = symmetrize(Sigma_y)

    # Cross-covariance: Sigma_x @ J^T
    cross_cov = Sigma @ J.T

    cov_y = lx.MatrixLinearOperator(Sigma_y, lx.positive_semidefinite_tag)
    out_state = GaussianState(mean=mu_y, cov=cov_y)

    return PropagationResult(state=out_state, cross_cov=cross_cov)

MonteCarloIntegrator

Bases: AbstractIntegrator

Monte Carlo moment matching: sample, propagate, compute moments.

Propagates uncertainty by drawing samples from the input Gaussian, evaluating the function at each sample, and computing empirical output moments:

x_i ~ N(mu, Sigma)       (n_samples points)
y_i = f(x_i)
mu_y = mean(y_i)
Sigma_y = cov(y_i) + regularization * I
cross_cov = cov(x_i, y_i)

Attributes:

Name	Type	Description
n_samples	int	Number of Monte Carlo samples. Default 1000.
regularization	float	Diagonal jitter for numerical stability.
key	Array | None	PRNG key. If None, uses jax.random.key(0).

Source code in src/gaussx/_quadrature/_monte_carlo.py

class MonteCarloIntegrator(AbstractIntegrator):
    r"""Monte Carlo moment matching: sample, propagate, compute moments.

    Propagates uncertainty by drawing samples from the input Gaussian,
    evaluating the function at each sample, and computing empirical
    output moments:

        x_i ~ N(mu, Sigma)       (n_samples points)
        y_i = f(x_i)
        mu_y = mean(y_i)
        Sigma_y = cov(y_i) + regularization * I
        cross_cov = cov(x_i, y_i)

    Attributes:
        n_samples: Number of Monte Carlo samples. Default ``1000``.
        regularization: Diagonal jitter for numerical stability.
        key: PRNG key. If ``None``, uses ``jax.random.key(0)``.
    """

    n_samples: int = eqx.field(static=True, default=1000)
    regularization: float = eqx.field(static=True, default=1e-6)
    key: jax.Array | None = None

    def guarantees_psd(self, dim: int) -> bool:
        """A sample covariance is PSD."""
        del dim
        return True

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via Monte Carlo sampling."""
        chi, w_m, w_c = self.points_and_weights(state)
        mu = state.mean

        # Propagate samples
        Y = jax.vmap(fn)(chi)  # (S, M)

        result = assemble_propagation_result(chi, Y, mu, w_m, w_c)

        # Add regularization jitter to output covariance
        Sigma_y = result.state.cov.as_matrix()
        M = Y.shape[1]
        import lineax as lx

        Sigma_y = Sigma_y + self.regularization * jnp.eye(M, dtype=Sigma_y.dtype)
        cov_y = lx.MatrixLinearOperator(Sigma_y, lx.positive_semidefinite_tag)
        out_state = GaussianState(mean=result.state.mean, cov=cov_y)

        return PropagationResult(state=out_state, cross_cov=result.cross_cov)

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """Return the Monte Carlo samples and their weights.

        Args:
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(points, w_m, w_c)`` with ``P = n_samples``. Mean
            weights are ``1/P``; covariance weights carry the Bessel
            correction ``1/(P − 1)``.

        Raises:
            ValueError: If ``n_samples < 2``.
        """
        if self.n_samples < 2:
            msg = (
                f"MonteCarloIntegrator requires n_samples >= 2 for "
                f"Bessel-corrected covariance, got {self.n_samples}."
            )
            raise ValueError(msg)

        mu = state.mean
        N = mu.shape[0]

        key = self.key if self.key is not None else jr.key(0)

        # Sample from input Gaussian: xᵢ = μ + L εᵢ. Draws and weights are
        # built in the input dtype: the defaults would otherwise land in
        # float64 under x64 and promote the whole transform.
        L = cholesky(state.cov).as_matrix()
        eps = jr.normal(key, (self.n_samples, N), dtype=mu.dtype)
        chi = mu[None, :] + eps @ L.T  # (S, N)

        # Uniform weights = 1/S for empirical moments
        # Use 1/(S−1) Bessel correction via covariance weights
        S = self.n_samples
        w_m = jnp.full(S, 1.0 / S, dtype=mu.dtype)
        w_c = jnp.full(S, 1.0 / (S - 1), dtype=mu.dtype)

        return chi, w_m, w_c

guarantees_psd(dim: int) -> bool

A sample covariance is PSD.

Source code in src/gaussx/_quadrature/_monte_carlo.py

def guarantees_psd(self, dim: int) -> bool:
    """A sample covariance is PSD."""
    del dim
    return True

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via Monte Carlo sampling.

Source code in src/gaussx/_quadrature/_monte_carlo.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via Monte Carlo sampling."""
    chi, w_m, w_c = self.points_and_weights(state)
    mu = state.mean

    # Propagate samples
    Y = jax.vmap(fn)(chi)  # (S, M)

    result = assemble_propagation_result(chi, Y, mu, w_m, w_c)

    # Add regularization jitter to output covariance
    Sigma_y = result.state.cov.as_matrix()
    M = Y.shape[1]
    import lineax as lx

    Sigma_y = Sigma_y + self.regularization * jnp.eye(M, dtype=Sigma_y.dtype)
    cov_y = lx.MatrixLinearOperator(Sigma_y, lx.positive_semidefinite_tag)
    out_state = GaussianState(mean=result.state.mean, cov=cov_y)

    return PropagationResult(state=out_state, cross_cov=result.cross_cov)

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Return the Monte Carlo samples and their weights.

Parameters:

Name	Type	Description	Default
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (points, w_m, w_c) with P = n_samples. Mean
Float[Array, ' P']	weights are 1/P; covariance weights carry the Bessel
Float[Array, ' P']	correction 1/(P − 1).

Raises:

Type	Description
ValueError	If n_samples < 2.

Source code in src/gaussx/_quadrature/_monte_carlo.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """Return the Monte Carlo samples and their weights.

    Args:
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(points, w_m, w_c)`` with ``P = n_samples``. Mean
        weights are ``1/P``; covariance weights carry the Bessel
        correction ``1/(P − 1)``.

    Raises:
        ValueError: If ``n_samples < 2``.
    """
    if self.n_samples < 2:
        msg = (
            f"MonteCarloIntegrator requires n_samples >= 2 for "
            f"Bessel-corrected covariance, got {self.n_samples}."
        )
        raise ValueError(msg)

    mu = state.mean
    N = mu.shape[0]

    key = self.key if self.key is not None else jr.key(0)

    # Sample from input Gaussian: xᵢ = μ + L εᵢ. Draws and weights are
    # built in the input dtype: the defaults would otherwise land in
    # float64 under x64 and promote the whole transform.
    L = cholesky(state.cov).as_matrix()
    eps = jr.normal(key, (self.n_samples, N), dtype=mu.dtype)
    chi = mu[None, :] + eps @ L.T  # (S, N)

    # Uniform weights = 1/S for empirical moments
    # Use 1/(S−1) Bessel correction via covariance weights
    S = self.n_samples
    w_m = jnp.full(S, 1.0 / S, dtype=mu.dtype)
    w_c = jnp.full(S, 1.0 / (S - 1), dtype=mu.dtype)

    return chi, w_m, w_c

AssumedDensityFilter

Bases: AbstractIntegrator

KL-optimal Gaussian projection via moment matching.

Projects the (possibly non-Gaussian) output distribution onto the Gaussian family by matching first and second moments. Equivalent to argmin_q KL(p(y) || q(y)) within the Gaussian family.

Adds adaptive regularization and optional diagnostics for detecting non-Gaussianity:

eps = eps_base * trace(Sigma_y) / n_dim

Attributes:

Name	Type	Description
n_samples	int	Number of Monte Carlo samples. Default 5000.
regularization	float	Base regularization. Default 1e-6.
adaptive_regularization	bool	Scale regularization by output variance. Default True.
key	Array | None	PRNG key. If None, uses jax.random.key(0).

Source code in src/gaussx/_quadrature/_adf.py

class AssumedDensityFilter(AbstractIntegrator):
    r"""KL-optimal Gaussian projection via moment matching.

    Projects the (possibly non-Gaussian) output distribution onto the
    Gaussian family by matching first and second moments. Equivalent to
    ``argmin_q KL(p(y) || q(y))`` within the Gaussian family.

    Adds adaptive regularization and optional diagnostics for detecting
    non-Gaussianity:

        eps = eps_base * trace(Sigma_y) / n_dim

    Attributes:
        n_samples: Number of Monte Carlo samples. Default ``5000``.
        regularization: Base regularization. Default ``1e-6``.
        adaptive_regularization: Scale regularization by output
            variance. Default ``True``.
        key: PRNG key. If ``None``, uses ``jax.random.key(0)``.
    """

    n_samples: int = eqx.field(static=True, default=5000)
    regularization: float = eqx.field(static=True, default=1e-6)
    adaptive_regularization: bool = eqx.field(static=True, default=True)
    key: jax.Array | None = None

    def guarantees_psd(self, dim: int) -> bool:
        """A regularised sample covariance is PSD."""
        del dim
        return True

    def points_and_weights(
        self,
        state: GaussianState,
    ) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
        """The `MonteCarloIntegrator` samples and weights for this rule.

        Lets `gaussx.moment_match` / `gaussx.statistical_linear_regression`
        use the same samples ``integrate`` does.

        Raises:
            ValueError: If ``n_samples < 2``.
        """
        return MonteCarloIntegrator(
            n_samples=self.n_samples, key=self.key
        ).points_and_weights(state)

    def integrate(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> PropagationResult:
        """Propagate Gaussian via assumed density filtering."""
        result, _ = self._integrate_impl(fn, state)
        return result

    def integrate_with_diagnostics(
        self,
        fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
        state: GaussianState,
    ) -> tuple[PropagationResult, dict]:
        """Propagate Gaussian and return non-Gaussianity diagnostics.

        Args:
            fn: Nonlinear function mapping ``(N,) -> (M,)``.
            state: Input Gaussian distribution.

        Returns:
            Tuple ``(result, diagnostics)`` where diagnostics contains
            ``skewness``, ``kurtosis``, ``min_eigval``, and
            ``condition_number``.
        """
        return self._integrate_impl(fn, state, compute_diagnostics=True)

    def _integrate_impl(
        self,
        fn: Callable,
        state: GaussianState,
        compute_diagnostics: bool = False,
    ) -> tuple[PropagationResult, dict]:
        """Core implementation with optional diagnostics."""
        # Same samples and Bessel-corrected moments as MonteCarloIntegrator
        # (shared points_and_weights + assemble_propagation_result).
        chi, w_m, w_c = self.points_and_weights(state)
        y_samples = jax.vmap(fn)(chi)
        moments = assemble_propagation_result(chi, y_samples, state.mean, w_m, w_c)
        mu_y = moments.state.mean
        Sigma_y = moments.state.cov.as_matrix()
        M = mu_y.shape[0]

        # Adaptive regularization
        if self.adaptive_regularization:
            eps_reg = self.regularization * jnp.trace(Sigma_y) / M
        else:
            eps_reg = self.regularization
        Sigma_y = symmetrize(Sigma_y + eps_reg * jnp.eye(M, dtype=Sigma_y.dtype))

        cov_y = lx.MatrixLinearOperator(Sigma_y, lx.positive_semidefinite_tag)
        out_state = GaussianState(mean=mu_y, cov=cov_y)
        result = PropagationResult(state=out_state, cross_cov=moments.cross_cov)

        diagnostics: dict = {}
        if compute_diagnostics:
            # Compute non-Gaussianity diagnostics via the gaussx
            # eigvals primitive (dispatches on operator structure).
            eigvals = _gaussx_eigvals(cov_y)
            min_eigval = jnp.min(eigvals)
            max_eigval = jnp.max(eigvals)
            cond = max_eigval / jnp.maximum(min_eigval, 1e-30)

            # Per-dimension skewness and kurtosis
            dy = einx.subtract("s m, m -> s m", y_samples, mu_y)
            std_dy = einx.divide("s m, m -> s m", dy, jnp.sqrt(jnp.diag(Sigma_y)))
            skewness = reduce(std_dy**3, "s m -> m", "mean")
            kurtosis = reduce(std_dy**4, "s m -> m", "mean")

            diagnostics = {
                "skewness": skewness,
                "kurtosis": kurtosis,
                "min_eigval": min_eigval,
                "condition_number": cond,
            }

        return result, diagnostics

guarantees_psd(dim: int) -> bool

A regularised sample covariance is PSD.

Source code in src/gaussx/_quadrature/_adf.py

def guarantees_psd(self, dim: int) -> bool:
    """A regularised sample covariance is PSD."""
    del dim
    return True

points_and_weights(state: GaussianState) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

The MonteCarloIntegrator samples and weights for this rule.

Lets gaussx.moment_match / gaussx.statistical_linear_regression use the same samples integrate does.

Raises:

Type	Description
ValueError	If n_samples < 2.

Source code in src/gaussx/_quadrature/_adf.py

def points_and_weights(
    self,
    state: GaussianState,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    """The `MonteCarloIntegrator` samples and weights for this rule.

    Lets `gaussx.moment_match` / `gaussx.statistical_linear_regression`
    use the same samples ``integrate`` does.

    Raises:
        ValueError: If ``n_samples < 2``.
    """
    return MonteCarloIntegrator(
        n_samples=self.n_samples, key=self.key
    ).points_and_weights(state)

integrate(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> PropagationResult

Propagate Gaussian via assumed density filtering.

Source code in src/gaussx/_quadrature/_adf.py

def integrate(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> PropagationResult:
    """Propagate Gaussian via assumed density filtering."""
    result, _ = self._integrate_impl(fn, state)
    return result

integrate_with_diagnostics(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState) -> tuple[PropagationResult, dict]

Propagate Gaussian and return non-Gaussianity diagnostics.

Parameters:

Name	Type	Description	Default
fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	Nonlinear function mapping (N,) -> (M,).	required
state	GaussianState	Input Gaussian distribution.	required

Returns:

Type	Description
PropagationResult	Tuple (result, diagnostics) where diagnostics contains
dict	skewness, kurtosis, min_eigval, and
tuple[PropagationResult, dict]	condition_number.

Source code in src/gaussx/_quadrature/_adf.py

def integrate_with_diagnostics(
    self,
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
) -> tuple[PropagationResult, dict]:
    """Propagate Gaussian and return non-Gaussianity diagnostics.

    Args:
        fn: Nonlinear function mapping ``(N,) -> (M,)``.
        state: Input Gaussian distribution.

    Returns:
        Tuple ``(result, diagnostics)`` where diagnostics contains
        ``skewness``, ``kurtosis``, ``min_eigval``, and
        ``condition_number``.
    """
    return self._integrate_impl(fn, state, compute_diagnostics=True)

moment_transform(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], mean: Float[Array, ' N'], cov: Float[Array, 'N N'], *, integrator: AbstractIntegrator | None = None) -> tuple[Float[Array, ' M'], Float[Array, 'M M'], Float[Array, 'N M']]

Moment triple of (x, fn(x)) for x Gaussian.

The Gaussian moment transform,

\[ \mathcal{T}[g; \mu, \Sigma] = (\mu_g,\ \Sigma_g,\ \Sigma_{xg}), \qquad x \sim \mathcal{N}(\mu, \Sigma), \]

with \(\mu_g = \mathbb{E}[g(x)]\), \(\Sigma_g = \mathrm{Cov}[g(x)]\) and \(\Sigma_{xg} = \mathrm{Cov}[x, g(x)]\).

This is the one method-specific step of every Gaussian filter: the filter loop, the smoother backward pass and the log-likelihood are identical across methods, and only \(\mathcal{T}\) changes. Taking dense arrays and returning dense arrays, it is the convenient form for building a filter — gaussx.nonlinear_kalman_filter is built on it — while AbstractIntegrator.integrate is the operator-typed form.

Unlike calling integrate directly, this requires the cross-covariance and says so if it is missing, rather than letting a None propagate into a silently zero Kalman gain.

Parameters:

Name	Type	Description	Default
fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	The map (N,) -> (M,) to push the Gaussian through.	required
mean	Float[Array, ' N']	Input mean \(\mu\), shape (N,).	required
cov	Float[Array, 'N N']	Input covariance \(\Sigma\), shape (N, N). Assumed PSD.	required
integrator	AbstractIntegrator | None	Rule realising \(\mathcal{T}\). Defaults to UnscentedIntegrator(alpha=1.0). Note this is not gaussx.UnscentedIntegrator's own alpha=1e-3 default, which is not safe in float32.	None

Returns:

Type	Description
Float[Array, ' M']	Tuple (mean_out, cov_out, cross_cov) with shapes (M,),
Float[Array, 'M M']	(M, M) and (N, M).

Raises:

Type	Description
TypeError	If the integrator does not supply a cross-covariance. Raised at trace time, so it costs nothing per step.

Source code in src/gaussx/_quadrature/_integrator.py

def moment_transform(
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    mean: Float[Array, " N"],
    cov: Float[Array, "N N"],
    *,
    integrator: AbstractIntegrator | None = None,
) -> tuple[Float[Array, " M"], Float[Array, "M M"], Float[Array, "N M"]]:
    r"""Moment triple of ``(x, fn(x))`` for ``x`` Gaussian.

    The Gaussian moment transform,

    $$
    \mathcal{T}[g; \mu, \Sigma] = (\mu_g,\ \Sigma_g,\ \Sigma_{xg}),
    \qquad x \sim \mathcal{N}(\mu, \Sigma),
    $$

    with $\mu_g = \mathbb{E}[g(x)]$, $\Sigma_g = \mathrm{Cov}[g(x)]$ and
    $\Sigma_{xg} = \mathrm{Cov}[x, g(x)]$.

    This is the one method-specific step of every Gaussian filter: the
    filter loop, the smoother backward pass and the log-likelihood are
    identical across methods, and only $\mathcal{T}$ changes. Taking
    dense arrays and returning dense arrays, it is the convenient form for
    building a filter — `gaussx.nonlinear_kalman_filter` is built on it —
    while `AbstractIntegrator.integrate` is the operator-typed form.

    Unlike calling ``integrate`` directly, this **requires** the
    cross-covariance and says so if it is missing, rather than letting a
    ``None`` propagate into a silently zero Kalman gain.

    Args:
        fn: The map ``(N,) -> (M,)`` to push the Gaussian through.
        mean: Input mean $\mu$, shape ``(N,)``.
        cov: Input covariance $\Sigma$, shape ``(N, N)``. Assumed PSD.
        integrator: Rule realising $\mathcal{T}$. Defaults to
            ``UnscentedIntegrator(alpha=1.0)``. Note this is *not*
            `gaussx.UnscentedIntegrator`'s own ``alpha=1e-3`` default,
            which is not safe in float32.

    Returns:
        Tuple ``(mean_out, cov_out, cross_cov)`` with shapes ``(M,)``,
        ``(M, M)`` and ``(N, M)``.

    Raises:
        TypeError: If the integrator does not supply a cross-covariance.
            Raised at trace time, so it costs nothing per step.
    """
    if integrator is None:
        # lazy import, cycle: _quadrature._unscented -> _quadrature._integrator
        from gaussx._quadrature._unscented import UnscentedIntegrator

        # alpha=1.0, not UnscentedIntegrator's own 1e-3: that spread
        # recovers even affine moments by cancellation between weights of
        # order 1e6, which costs ~7 digits and is ruinous in float32.
        integrator = UnscentedIntegrator(alpha=1.0)

    state = GaussianState(
        mean=mean, cov=lx.MatrixLinearOperator(cov, lx.positive_semidefinite_tag)
    )
    result = integrator.integrate(fn, state)
    if result.cross_cov is None:
        msg = (
            f"{type(integrator).__name__} returned cross_cov=None, but the "
            f"moment transform is defined by all three moments. Downstream "
            f"consumers build a Kalman gain as cross_cov @ S^-1, so a "
            f"missing cross-covariance would silently give a zero gain."
        )
        raise TypeError(msg)
    return result.state.mean, result.state.cov.as_matrix(), result.cross_cov

The integrators are also what selects the filter in nonlinear_kalman_filter — EKF, UKF, CKF, GHKF and the Monte-Carlo variant are the same loop under a different rule.

Quadrature rules

The raw point sets behind the integrators, for when a recipe needs direct control.

Structured linear algebra and Gaussian primitives for JAX.

gauss_hermite_points(order: int, dim: int, dtype: DTypeLike | None = None) -> tuple[Float[Array, 'P D'], Float[Array, ' P']]

Gauss-Hermite quadrature points and weights.

Generates tensor-product Gauss-Hermite quadrature points for integrating functions against a standard Gaussian measure (probabilists' Hermite polynomials). P = order^dim total points.

Parameters:

Name	Type	Description	Default
order	int	Number of quadrature points per dimension.	required
dim	int	Dimensionality of the integration domain.	required
dtype	DTypeLike | None	Floating dtype of the returned points and weights. Defaults to JAX's default float dtype. Unlike the other rules here this one has no mean to take a dtype from, so callers propagating an input dtype must pass it explicitly — gaussx.GaussHermiteIntegrator does.	None

Returns:

Type	Description
Float[Array, 'P D']	Tuple (points, weights) where points has shape
Float[Array, ' P']	(order^dim, dim) and weights has shape (order^dim,).

Source code in src/gaussx/_quadrature/_quadrature.py

def gauss_hermite_points(
    order: int,
    dim: int,
    dtype: DTypeLike | None = None,
) -> tuple[Float[Array, "P D"], Float[Array, " P"]]:
    r"""Gauss-Hermite quadrature points and weights.

    Generates tensor-product Gauss-Hermite quadrature points for
    integrating functions against a standard Gaussian measure
    (probabilists' Hermite polynomials).
    ``P = order^dim`` total points.

    Args:
        order: Number of quadrature points per dimension.
        dim: Dimensionality of the integration domain.
        dtype: Floating dtype of the returned points and weights.
            Defaults to JAX's default float dtype. Unlike the other rules
            here this one has no mean to take a dtype from, so callers
            propagating an input dtype must pass it explicitly —
            `gaussx.GaussHermiteIntegrator` does.

    Returns:
        Tuple ``(points, weights)`` where points has shape
        ``(order^dim, dim)`` and weights has shape ``(order^dim,)``.
    """
    # 1D probabilists' Gauss-Hermite (weight = exp(-x^2/2))
    x1d_np, w1d_np = np.polynomial.hermite_e.hermegauss(order)
    x1d = jnp.array(x1d_np, dtype=dtype)
    w1d = jnp.array(w1d_np, dtype=dtype)

    if dim == 1:
        return x1d[:, None], w1d

    # Tensor product via meshgrid
    grids = jnp.meshgrid(*([x1d] * dim), indexing="ij")
    stacked = jnp.stack(grids, axis=0)  # (dim, *grid_shape)
    points = rearrange(stacked, "D ... -> (...) D")

    weight_grids = jnp.meshgrid(*([w1d] * dim), indexing="ij")
    weight_stack = jnp.stack(weight_grids, axis=0)  # (dim, *grid_shape)
    weights = reduce(weight_stack, "D ... -> (...)", "prod")

    return points, weights

cubature_points(mean: Float[Array, ' N'], cov: lx.AbstractLinearOperator) -> tuple[Float[Array, 'P N'], Float[Array, ' P']]

Spherical-radial cubature points and weights.

Generates 2N cubature points with equal weights 1/(2N). This is the cubature Kalman filter (CKF) point set.

Uses gaussx.sqrt(cov) for structured square root dispatch.

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' N']	Mean vector, shape (N,).	required
cov	AbstractLinearOperator	Covariance operator, shape (N, N).	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (chi, weights) where:
Float[Array, ' P']	chi: Cubature points, shape (2N, N).
tuple[Float[Array, 'P N'], Float[Array, ' P']]	weights: Equal weights, shape (2N,).

Source code in src/gaussx/_quadrature/_quadrature.py

def cubature_points(
    mean: Float[Array, " N"],
    cov: lx.AbstractLinearOperator,
) -> tuple[Float[Array, "P N"], Float[Array, " P"]]:
    r"""Spherical-radial cubature points and weights.

    Generates ``2N`` cubature points with equal weights ``1/(2N)``.
    This is the cubature Kalman filter (CKF) point set.

    Uses ``gaussx.sqrt(cov)`` for structured square root dispatch.

    Args:
        mean: Mean vector, shape ``(N,)``.
        cov: Covariance operator, shape ``(N, N)``.

    Returns:
        Tuple ``(chi, weights)`` where:
        - ``chi``: Cubature points, shape ``(2N, N)``.
        - ``weights``: Equal weights, shape ``(2N,)``.
    """
    N = mean.shape[0]

    S = sqrt(cov)
    S_mat = S.as_matrix()
    S_scaled = jnp.sqrt(N) * S_mat  # (N, N)

    chi_plus = mean[None, :] + S_scaled.T  # (N, N)
    chi_minus = mean[None, :] - S_scaled.T  # (N, N)
    chi = jnp.concatenate([chi_plus, chi_minus], axis=0)  # (2N, N)

    weights = jnp.full(2 * N, 1.0 / (2.0 * N), dtype=mean.dtype)

    return chi, weights

fifth_order_cubature_points(mean: Float[Array, ' N'], cov: lx.AbstractLinearOperator) -> tuple[Float[Array, 'P N'], Float[Array, ' P']]

Fifth-order fully symmetric cubature points and weights.

The McNamee & Stenger (1967) degree-5 rule: exact for polynomials up to total degree 5 in N variables, using P = 2 N^2 + 1 points laid out on two shells plus the origin,

========== =========================== ============== ==================== Shell Points Multiplicity Weight ========== =========================== ============== ==================== origin 0 1 2 / (N + 2) axis ±r₁ eᵢ 2N (4 − N) / (2 (N + 2)²) diagonal ±r₂ (eᵢ ± eⱼ), i < j 2N(N − 1) 1 / (N + 2)² ========== =========================== ============== ====================

with r₁ = sqrt(N + 2) and r₂ = sqrt((N + 2) / 2). Points are mapped into the input space as χᵢ = mean + S zᵢ where S = gaussx.sqrt(cov).

This rule sits between sigma_points / cubature_points (2N + 1 and 2N points, degree 3) and a tensor-product gauss_hermite_points rule (order^N points), and is the default integrator in BayesNewton.

Note

The axis weight is negative for N > 4 (and zero at N = 4). That is inherent to the rule, not a bug: the weights still sum to one and every moment through degree 5 is reproduced exactly, but the rule is not a positive quadrature in high dimension.

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' N']	Mean vector, shape (N,).	required
cov	AbstractLinearOperator	Covariance operator, shape (N, N).	required

Returns:

Type	Description
Float[Array, 'P N']	Tuple (chi, weights) where:
Float[Array, ' P']	chi: Cubature points, shape (2 N^2 + 1, N).
tuple[Float[Array, 'P N'], Float[Array, ' P']]	weights: Weights summing to one, shape (2 N^2 + 1,).

Source code in src/gaussx/_quadrature/_quadrature.py

def fifth_order_cubature_points(
    mean: Float[Array, " N"],
    cov: lx.AbstractLinearOperator,
) -> tuple[Float[Array, "P N"], Float[Array, " P"]]:
    r"""Fifth-order fully symmetric cubature points and weights.

    The McNamee & Stenger (1967) degree-5 rule: exact for polynomials up
    to total degree 5 in ``N`` variables, using ``P = 2 N^2 + 1`` points
    laid out on two shells plus the origin,

    ==========  ===========================  ==============  ====================
    Shell       Points                       Multiplicity    Weight
    ==========  ===========================  ==============  ====================
    origin      ``0``                        ``1``           ``2 / (N + 2)``
    axis        ``±r₁ eᵢ``                   ``2N``          ``(4 − N) / (2 (N + 2)²)``
    diagonal    ``±r₂ (eᵢ ± eⱼ)``, ``i < j`` ``2N(N − 1)``   ``1 / (N + 2)²``
    ==========  ===========================  ==============  ====================

    with ``r₁ = sqrt(N + 2)`` and ``r₂ = sqrt((N + 2) / 2)``. Points are
    mapped into the input space as ``χᵢ = mean + S zᵢ`` where
    ``S = gaussx.sqrt(cov)``.

    This rule sits between `sigma_points` / `cubature_points` (``2N + 1``
    and ``2N`` points, degree 3) and a tensor-product `gauss_hermite_points`
    rule (``order^N`` points), and is the default integrator in BayesNewton.

    Note:
        The axis weight is negative for ``N > 4`` (and zero at ``N = 4``).
        That is inherent to the rule, not a bug: the weights still sum to
        one and every moment through degree 5 is reproduced exactly, but
        the rule is not a positive quadrature in high dimension.

    Args:
        mean: Mean vector, shape ``(N,)``.
        cov: Covariance operator, shape ``(N, N)``.

    Returns:
        Tuple ``(chi, weights)`` where:
        - ``chi``: Cubature points, shape ``(2 N^2 + 1, N)``.
        - ``weights``: Weights summing to one, shape ``(2 N^2 + 1,)``.
    """
    N = mean.shape[0]

    # Python floats, not numpy scalars: a np.float64 radius is strongly
    # typed and would promote the points it scales back to float64.
    r1 = float(np.sqrt(N + 2.0))
    r2 = float(np.sqrt((N + 2.0) / 2.0))
    w0 = 2.0 / (N + 2.0)
    w1 = (4.0 - N) / (2.0 * (N + 2.0) ** 2)
    w2 = 1.0 / (N + 2.0) ** 2

    # Points and weights are built in the input dtype: the defaults would
    # otherwise land in float64 under x64 and promote the whole transform.
    dtype = mean.dtype

    eye = jnp.eye(N, dtype=dtype)
    origin = jnp.zeros((1, N), dtype=dtype)
    axis = jnp.concatenate([r1 * eye, -r1 * eye], axis=0)  # (2N, N)

    # Diagonal shell: all four sign combinations of ±eᵢ ± eⱼ for i < j.
    rows, cols = np.triu_indices(N, k=1)
    n_pairs = rows.size
    shells = [origin, axis]
    if n_pairs:
        e_i = eye[rows]  # (n_pairs, N)
        e_j = eye[cols]
        shells.append(
            jnp.concatenate(
                [
                    r2 * (e_i + e_j),
                    r2 * (e_i - e_j),
                    -r2 * (e_i + e_j),
                    -r2 * (e_i - e_j),
                ],
                axis=0,
            )
        )
    z = jnp.concatenate(shells, axis=0)  # (2N² + 1, N)

    weights = jnp.concatenate(
        [
            jnp.array([w0], dtype=dtype),
            jnp.full(2 * N, w1, dtype=dtype),
            jnp.full(4 * n_pairs, w2, dtype=dtype),
        ]
    )

    # Map into the input space: χᵢ = μ + S zᵢ with S Sᵀ = cov.
    S = sqrt(cov).as_matrix()
    chi = mean[None, :] + z @ S.T

    return chi, weights

sigma_points(mean: Float[Array, ' N'], cov: lx.AbstractLinearOperator, alpha: float = 0.001, beta: float = 2.0, kappa: float = 0.0) -> tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]

Unscented transform sigma points and weights.

Generates 2N+1 deterministic sigma points for a Gaussian with the given mean and covariance, using the scaled unscented transform.

Uses gaussx.sqrt(cov) for structured square root dispatch.

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' N']	Mean vector, shape (N,).	required
cov	AbstractLinearOperator	Covariance operator, shape (N, N).	required
alpha	float	Spread parameter. Controls how far sigma points are from the mean. Default 1e-3 (the classic Wan-van der Merwe scaled transform). See the warning below; alpha=1.0 is the recommended choice, especially in float32.	0.001
beta	float	Prior distribution parameter. beta=2 is optimal for Gaussians. Default 2.0.	2.0
kappa	float	Secondary scaling parameter. Default 0.0.	0.0

Warning

With lambda = alpha^2 (N + kappa) - N the centre mean weight is lambda / (N + lambda) \approx 1 - 1/alpha^2: about -1e6 at the default alpha=1e-3. The moments are then recovered by cancellation between weights of order 1e6, which costs about six digits; in float32 a smooth nonlinearity comes out wrong in the second digit (gh-310). alpha=1.0 (with beta=2, kappa=0) is the symmetric 2N+1 rule with a zero centre mean weight, exact for affine maps, and is what every gaussx-internal default (moment_transform, the nonlinear Kalman filters) uses. A UserWarning is emitted when the actual centre mean weight (computed from alpha, kappa and N) has magnitude >= 1e3 with a float32 (or lower precision) mean; it is skipped when alpha or kappa is traced.

Returns:

Type	Description
Float[Array, 'P N']	Tuple (chi, w_m, w_c) where:
Float[Array, ' P']	chi: Sigma points, shape (2N+1, N).
Float[Array, ' P']	w_m: Mean weights, shape (2N+1,).
tuple[Float[Array, 'P N'], Float[Array, ' P'], Float[Array, ' P']]	w_c: Covariance weights, shape (2N+1,).

Source code in src/gaussx/_quadrature/_quadrature.py

def sigma_points(
    mean: Float[Array, " N"],
    cov: lx.AbstractLinearOperator,
    alpha: float = 1e-3,
    beta: float = 2.0,
    kappa: float = 0.0,
) -> tuple[Float[Array, "P N"], Float[Array, " P"], Float[Array, " P"]]:
    r"""Unscented transform sigma points and weights.

    Generates ``2N+1`` deterministic sigma points for a Gaussian with
    the given mean and covariance, using the scaled unscented transform.

    Uses ``gaussx.sqrt(cov)`` for structured square root dispatch.

    Args:
        mean: Mean vector, shape ``(N,)``.
        cov: Covariance operator, shape ``(N, N)``.
        alpha: Spread parameter. Controls how far sigma points are
            from the mean. Default ``1e-3`` (the classic Wan-van der Merwe
            scaled transform). See the warning below; ``alpha=1.0`` is the
            recommended choice, especially in float32.
        beta: Prior distribution parameter. ``beta=2`` is optimal for
            Gaussians. Default ``2.0``.
        kappa: Secondary scaling parameter. Default ``0.0``.

    Warning:
        With ``lambda = alpha^2 (N + kappa) - N`` the centre mean weight is
        ``lambda / (N + lambda) \approx 1 - 1/alpha^2``: about ``-1e6`` at
        the default ``alpha=1e-3``. The moments are then recovered by
        cancellation between weights of order ``1e6``, which costs about
        six digits; in float32 a smooth nonlinearity comes out wrong in
        the second digit (gh-310). ``alpha=1.0`` (with ``beta=2``,
        ``kappa=0``) is the symmetric ``2N+1`` rule with a zero centre mean
        weight, exact for affine maps, and is what every gaussx-internal
        default (`moment_transform`, the nonlinear Kalman filters) uses.
        A `UserWarning` is emitted when the actual centre mean weight
        (computed from ``alpha``, ``kappa`` and ``N``) has magnitude
        ``>= 1e3`` with a float32 (or lower precision) mean; it is skipped
        when ``alpha`` or ``kappa`` is traced.

    Returns:
        Tuple ``(chi, w_m, w_c)`` where:
        - ``chi``: Sigma points, shape ``(2N+1, N)``.
        - ``w_m``: Mean weights, shape ``(2N+1,)``.
        - ``w_c``: Covariance weights, shape ``(2N+1,)``.
    """
    _warn_cancelling_centre_weight(mean, alpha, kappa)
    N = mean.shape[0]
    lam = alpha**2 * (N + kappa) - N
    c = N + lam

    # Matrix square root: S where cov = S S^T
    S = sqrt(cov)
    S_mat = S.as_matrix()  # (N, N)
    S_scaled = jnp.sqrt(c) * S_mat

    # Sigma points: chi_0 = mu, chi_i = mu + S_i, chi_{N+i} = mu - S_i
    chi_0 = mean[None, :]  # (1, N)
    chi_plus = mean[None, :] + S_scaled.T  # (N, N) — each row is a point
    chi_minus = mean[None, :] - S_scaled.T  # (N, N)
    chi = jnp.concatenate([chi_0, chi_plus, chi_minus], axis=0)  # (2N+1, N)

    # Weights are built in the input dtype: Python floats would otherwise
    # land in float64 under x64 and promote the whole downstream transform.
    dtype = mean.dtype

    # Mean weights
    w_m_0 = lam / c
    w_m_i = 1.0 / (2.0 * c)
    w_m = jnp.concatenate(
        [
            jnp.array([w_m_0], dtype=dtype),
            jnp.full(2 * N, w_m_i, dtype=dtype),
        ]
    )

    # Covariance weights
    w_c_0 = lam / c + (1.0 - alpha**2 + beta)
    w_c = jnp.concatenate(
        [
            jnp.array([w_c_0], dtype=dtype),
            jnp.full(2 * N, w_m_i, dtype=dtype),
        ]
    )

    return chi, w_m, w_c

Likelihoods

Observation models with quadrature-friendly log_prob surfaces, shared by the expectation helpers and the SSM / CVI recipes. Site-factorising likelihoods also give site_derivatives, the per-site gradient and Hessian diagonal that Newton-type inference (newton_update, laplace_mode) consumes; binomial and negative binomial supply closed forms.

Structured linear algebra and Gaussian primitives for JAX.

AbstractLikelihood

Bases: Module

Base class for likelihood functions with optional analytical ELL.

Subclasses that support closed-form expected log-likelihood under a Gaussian variational distribution should override has_analytical_ell to return True and implement analytical_expected_log_likelihood.

Source code in src/gaussx/_quadrature/_likelihood.py

class AbstractLikelihood(eqx.Module):
    """Base class for likelihood functions with optional analytical ELL.

    Subclasses that support closed-form expected log-likelihood under
    a Gaussian variational distribution should override
    ``has_analytical_ell`` to return ``True`` and implement
    ``analytical_expected_log_likelihood``.
    """

    @abc.abstractmethod
    def log_prob(
        self,
        f: Float[Array, " N"],
    ) -> Float[Array, ""]:
        """Evaluate ``log p(y | f)`` for fixed observations.

        Args:
            f: Latent function values, shape ``(N,)``.

        Returns:
            Scalar log-likelihood.
        """
        ...

    def site_derivatives(
        self,
        f: Float[Array, " N"],
    ) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
        r"""Per-site gradient and Hessian diagonal of ``log p(y | f)`` in ``f``.

        For a likelihood that factorises over sites,
        $\log p(y\mid f) = \sum_i \log p(y_i\mid f_i)$, the Hessian is
        diagonal and this returns ``(∂ log p / ∂f_i, ∂² log p / ∂f_i²)``, the
        inputs of `gaussx.newton_update`'s diagonal path and of
        `gaussx.laplace_mode`. The default differentiates ``log_prob``: one
        reverse pass for the gradient and one forward pass of it along
        ``1`` for the diagonal (exact only when the Hessian is diagonal, so
        not for `SoftmaxLikelihood` or `HeteroscedasticGaussianLikelihood`).
        Subclasses may override it with closed forms.

        Args:
            f: Latent function values, shape ``(N,)``.

        Returns:
            Tuple ``(gradient, hessian_diagonal)``, each shape ``(N,)``.
        """
        return jax.jvp(jax.grad(self.log_prob), (f,), (jnp.ones_like(f),))

    def has_analytical_ell(self) -> bool:
        """Whether this likelihood supports closed-form ELL."""
        return False

    def analytical_expected_log_likelihood(
        self,
        q_mu: Float[Array, " N"],
        q_cov: lx.AbstractLinearOperator,
    ) -> Float[Array, ""]:
        """Closed-form ``E_q[log p(y | f)]`` where ``q = N(q_mu, q_cov)``.

        Args:
            q_mu: Variational mean, shape ``(N,)``.
            q_cov: Variational covariance operator, shape ``(N, N)``.

        Returns:
            Scalar expected log-likelihood.

        Raises:
            NotImplementedError: If no analytical form exists.
        """
        msg = f"{type(self).__name__} has no analytical ELL"
        raise NotImplementedError(msg)

log_prob(f: Float[Array, ' N']) -> Float[Array, ''] abstractmethod

Evaluate log p(y | f) for fixed observations.

Parameters:

Name	Type	Description	Default
f	Float[Array, ' N']	Latent function values, shape (N,).	required

Returns:

Type	Description
Float[Array, '']	Scalar log-likelihood.

Source code in src/gaussx/_quadrature/_likelihood.py

@abc.abstractmethod
def log_prob(
    self,
    f: Float[Array, " N"],
) -> Float[Array, ""]:
    """Evaluate ``log p(y | f)`` for fixed observations.

    Args:
        f: Latent function values, shape ``(N,)``.

    Returns:
        Scalar log-likelihood.
    """
    ...

site_derivatives(f: Float[Array, ' N']) -> tuple[Float[Array, ' N'], Float[Array, ' N']]

Per-site gradient and Hessian diagonal of log p(y | f) in f.

For a likelihood that factorises over sites, \(\log p(y\mid f) = \sum_i \log p(y_i\mid f_i)\), the Hessian is diagonal and this returns (∂ log p / ∂f_i, ∂² log p / ∂f_i²), the inputs of gaussx.newton_update's diagonal path and of gaussx.laplace_mode. The default differentiates log_prob: one reverse pass for the gradient and one forward pass of it along 1 for the diagonal (exact only when the Hessian is diagonal, so not for SoftmaxLikelihood or HeteroscedasticGaussianLikelihood). Subclasses may override it with closed forms.

Parameters:

Name	Type	Description	Default
f	Float[Array, ' N']	Latent function values, shape (N,).	required

Returns:

Type	Description
tuple[Float[Array, ' N'], Float[Array, ' N']]	Tuple (gradient, hessian_diagonal), each shape (N,).

Source code in src/gaussx/_quadrature/_likelihood.py

def site_derivatives(
    self,
    f: Float[Array, " N"],
) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
    r"""Per-site gradient and Hessian diagonal of ``log p(y | f)`` in ``f``.

    For a likelihood that factorises over sites,
    $\log p(y\mid f) = \sum_i \log p(y_i\mid f_i)$, the Hessian is
    diagonal and this returns ``(∂ log p / ∂f_i, ∂² log p / ∂f_i²)``, the
    inputs of `gaussx.newton_update`'s diagonal path and of
    `gaussx.laplace_mode`. The default differentiates ``log_prob``: one
    reverse pass for the gradient and one forward pass of it along
    ``1`` for the diagonal (exact only when the Hessian is diagonal, so
    not for `SoftmaxLikelihood` or `HeteroscedasticGaussianLikelihood`).
    Subclasses may override it with closed forms.

    Args:
        f: Latent function values, shape ``(N,)``.

    Returns:
        Tuple ``(gradient, hessian_diagonal)``, each shape ``(N,)``.
    """
    return jax.jvp(jax.grad(self.log_prob), (f,), (jnp.ones_like(f),))

has_analytical_ell() -> bool

Whether this likelihood supports closed-form ELL.

Source code in src/gaussx/_quadrature/_likelihood.py

def has_analytical_ell(self) -> bool:
    """Whether this likelihood supports closed-form ELL."""
    return False

analytical_expected_log_likelihood(q_mu: Float[Array, ' N'], q_cov: lx.AbstractLinearOperator) -> Float[Array, '']

Closed-form E_q[log p(y | f)] where q = N(q_mu, q_cov).

Parameters:

Name	Type	Description	Default
q_mu	Float[Array, ' N']	Variational mean, shape (N,).	required
q_cov	AbstractLinearOperator	Variational covariance operator, shape (N, N).	required

Returns:

Type	Description
Float[Array, '']	Scalar expected log-likelihood.

Raises:

Type	Description
NotImplementedError	If no analytical form exists.

Source code in src/gaussx/_quadrature/_likelihood.py

def analytical_expected_log_likelihood(
    self,
    q_mu: Float[Array, " N"],
    q_cov: lx.AbstractLinearOperator,
) -> Float[Array, ""]:
    """Closed-form ``E_q[log p(y | f)]`` where ``q = N(q_mu, q_cov)``.

    Args:
        q_mu: Variational mean, shape ``(N,)``.
        q_cov: Variational covariance operator, shape ``(N, N)``.

    Returns:
        Scalar expected log-likelihood.

    Raises:
        NotImplementedError: If no analytical form exists.
    """
    msg = f"{type(self).__name__} has no analytical ELL"
    raise NotImplementedError(msg)

GaussianLikelihood

Bases: AbstractLikelihood

Gaussian likelihood log N(y | f, noise_var * I).

Supports closed-form expected log-likelihood:

E_q[log N(y | f, \sigma^2 I)]
    = log N(y | q_\mu, \sigma^2 I)
      - 0.5 / \sigma^2 \cdot tr(q_{cov})

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Observed targets, shape (N,).
noise_var	float	Observation noise variance (scalar).

Source code in src/gaussx/_quadrature/_likelihood.py

class GaussianLikelihood(AbstractLikelihood):
    r"""Gaussian likelihood ``log N(y | f, noise_var * I)``.

    Supports closed-form expected log-likelihood:

        E_q[log N(y | f, \sigma^2 I)]
            = log N(y | q_\mu, \sigma^2 I)
              - 0.5 / \sigma^2 \cdot tr(q_{cov})

    Attributes:
        y: Observed targets, shape ``(N,)``.
        noise_var: Observation noise variance (scalar).
    """

    y: Float[Array, " N"]
    noise_var: float

    def log_prob(
        self,
        f: Float[Array, " N"],
    ) -> Float[Array, ""]:
        """Evaluate ``log N(y | f, noise_var * I)``."""
        N = self.y.shape[-1]
        residual = self.y - f
        return -0.5 * (
            N * _LOG_2PI
            + N * jnp.log(self.noise_var)
            + jnp.sum(residual**2) / self.noise_var
        )

    def has_analytical_ell(self) -> bool:
        """Gaussian likelihood has closed-form ELL."""
        return True

    def analytical_expected_log_likelihood(
        self,
        q_mu: Float[Array, " N"],
        q_cov: lx.AbstractLinearOperator,
    ) -> Float[Array, ""]:
        r"""Closed-form ``E_q[log N(y | f, \sigma^2 I)]``.

        Uses:

            E_q[log N(y|f,R)] = log N(y | q_mu, R) - 0.5 tr(R^{-1} q_cov)

        where ``R = noise_var * I``. Delegates the log-density term to
        `gaussx.gaussian_log_prob` (which exploits the diagonal
        noise structure) and computes the trace correction directly via
        the structural ``trace(q_cov) / noise_var`` shortcut, so
        Kronecker/BlockDiag-structured ``q_cov`` keeps its O(n)
        ``prod(trace_factor)`` / per-block ``trace`` fast paths instead
        of materializing through ``trace_product(R^{-1}, q_cov)``.
        """

        N = self.y.shape[-1]
        noise = lx.DiagonalLinearOperator(jnp.full(N, self.noise_var))
        log_pdf = gaussian_log_prob(q_mu, noise, self.y)
        # Structural fast path: tr(R^{-1} q_cov) = trace(q_cov) / noise_var
        # for scalar isotropic noise. ``trace`` dispatches on operator
        # structure (Kronecker, BlockDiag, …) via gaussx primitives.
        tr_term = trace(q_cov) / self.noise_var
        return log_pdf - 0.5 * tr_term

log_prob(f: Float[Array, ' N']) -> Float[Array, '']

Evaluate log N(y | f, noise_var * I).

Source code in src/gaussx/_quadrature/_likelihood.py

def log_prob(
    self,
    f: Float[Array, " N"],
) -> Float[Array, ""]:
    """Evaluate ``log N(y | f, noise_var * I)``."""
    N = self.y.shape[-1]
    residual = self.y - f
    return -0.5 * (
        N * _LOG_2PI
        + N * jnp.log(self.noise_var)
        + jnp.sum(residual**2) / self.noise_var
    )

has_analytical_ell() -> bool

Gaussian likelihood has closed-form ELL.

Source code in src/gaussx/_quadrature/_likelihood.py

def has_analytical_ell(self) -> bool:
    """Gaussian likelihood has closed-form ELL."""
    return True

analytical_expected_log_likelihood(q_mu: Float[Array, ' N'], q_cov: lx.AbstractLinearOperator) -> Float[Array, '']

Closed-form E_q[log N(y | f, \sigma^2 I)].

Uses:

E_q[log N(y|f,R)] = log N(y | q_mu, R) - 0.5 tr(R^{-1} q_cov)

where R = noise_var * I. Delegates the log-density term to gaussx.gaussian_log_prob (which exploits the diagonal noise structure) and computes the trace correction directly via the structural trace(q_cov) / noise_var shortcut, so Kronecker/BlockDiag-structured q_cov keeps its O(n) prod(trace_factor) / per-block trace fast paths instead of materializing through trace_product(R^{-1}, q_cov).

Source code in src/gaussx/_quadrature/_likelihood.py

def analytical_expected_log_likelihood(
    self,
    q_mu: Float[Array, " N"],
    q_cov: lx.AbstractLinearOperator,
) -> Float[Array, ""]:
    r"""Closed-form ``E_q[log N(y | f, \sigma^2 I)]``.

    Uses:

        E_q[log N(y|f,R)] = log N(y | q_mu, R) - 0.5 tr(R^{-1} q_cov)

    where ``R = noise_var * I``. Delegates the log-density term to
    `gaussx.gaussian_log_prob` (which exploits the diagonal
    noise structure) and computes the trace correction directly via
    the structural ``trace(q_cov) / noise_var`` shortcut, so
    Kronecker/BlockDiag-structured ``q_cov`` keeps its O(n)
    ``prod(trace_factor)`` / per-block ``trace`` fast paths instead
    of materializing through ``trace_product(R^{-1}, q_cov)``.
    """

    N = self.y.shape[-1]
    noise = lx.DiagonalLinearOperator(jnp.full(N, self.noise_var))
    log_pdf = gaussian_log_prob(q_mu, noise, self.y)
    # Structural fast path: tr(R^{-1} q_cov) = trace(q_cov) / noise_var
    # for scalar isotropic noise. ``trace`` dispatches on operator
    # structure (Kronecker, BlockDiag, …) via gaussx primitives.
    tr_term = trace(q_cov) / self.noise_var
    return log_pdf - 0.5 * tr_term

HeteroscedasticGaussianLikelihood

Bases: AbstractLikelihood

Heteroscedastic Gaussian likelihood with input-dependent noise.

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Observations, shape (N,).

Source code in src/gaussx/_quadrature/_likelihoods.py

class HeteroscedasticGaussianLikelihood(AbstractLikelihood):
    r"""Heteroscedastic Gaussian likelihood with input-dependent noise.

    Attributes:
        y: Observations, shape ``(N,)``.
    """

    y: Float[Array, " N"]
    latent_dim: int = eqx.field(static=True, default=2)

    def log_prob(self, f: Float[Array, " 2N"]) -> Float[Array, ""]:
        """Evaluate heteroscedastic Gaussian log-likelihood."""
        N = self.y.shape[0]
        f_mean = f[:N]
        f_noise = f[N:]
        noise_std = jax.nn.softplus(f_noise)
        noise_var = noise_std**2

        log_2pi = jnp.log(2.0 * jnp.pi)
        residual = self.y - f_mean
        return jnp.sum(-0.5 * (log_2pi + jnp.log(noise_var) + residual**2 / noise_var))

log_prob(f: Float[Array, ' 2N']) -> Float[Array, '']

Evaluate heteroscedastic Gaussian log-likelihood.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " 2N"]) -> Float[Array, ""]:
    """Evaluate heteroscedastic Gaussian log-likelihood."""
    N = self.y.shape[0]
    f_mean = f[:N]
    f_noise = f[N:]
    noise_std = jax.nn.softplus(f_noise)
    noise_var = noise_std**2

    log_2pi = jnp.log(2.0 * jnp.pi)
    residual = self.y - f_mean
    return jnp.sum(-0.5 * (log_2pi + jnp.log(noise_var) + residual**2 / noise_var))

BernoulliLikelihood

Bases: AbstractLikelihood

Bernoulli likelihood with logit link.

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Binary observations, shape (N,).

Source code in src/gaussx/_quadrature/_likelihoods.py

class BernoulliLikelihood(AbstractLikelihood):
    r"""Bernoulli likelihood with logit link.

    Attributes:
        y: Binary observations, shape ``(N,)``.
    """

    y: Float[Array, " N"]

    def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
        """Evaluate Bernoulli log-likelihood with logit link."""
        return jnp.sum(
            self.y * jax.nn.log_sigmoid(f) + (1.0 - self.y) * jax.nn.log_sigmoid(-f)
        )

log_prob(f: Float[Array, ' N']) -> Float[Array, '']

Evaluate Bernoulli log-likelihood with logit link.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
    """Evaluate Bernoulli log-likelihood with logit link."""
    return jnp.sum(
        self.y * jax.nn.log_sigmoid(f) + (1.0 - self.y) * jax.nn.log_sigmoid(-f)
    )

BinomialLikelihood

Bases: AbstractLikelihood

Binomial likelihood with logit link.

\(y_i \sim \operatorname{Bin}(n_i, \sigma(f_i))\), so

\[ \log p(y\mid f) = \sum_i \log\binom{n_i}{y_i} + y_i\log\sigma(f_i) + (n_i - y_i)\log\sigma(-f_i), \]

with site gradient \(y_i - n_i\sigma(f_i)\) and Hessian \(-n_i\sigma(f_i)\sigma(-f_i)\) (log-concave). n_trials = 1 is BernoulliLikelihood.

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Success counts, shape (N,).
n_trials	Float[Array, ' N'] | float	Numbers of trials \(n_i\), a scalar or shape (N,).

Examples:

import jax.numpy as jnp
import gaussx as gx

lik = gx.BinomialLikelihood(jnp.array([3.0, 0.0, 5.0]), n_trials=5.0)
lp = lik.log_prob(jnp.zeros(3))
grad, hess = lik.site_derivatives(jnp.zeros(3))  # [0.5, -2.5, 2.5], -1.25

Source code in src/gaussx/_quadrature/_likelihoods.py

class BinomialLikelihood(AbstractLikelihood):
    r"""Binomial likelihood with logit link.

    $y_i \sim \operatorname{Bin}(n_i, \sigma(f_i))$, so

    $$
    \log p(y\mid f) = \sum_i \log\binom{n_i}{y_i} + y_i\log\sigma(f_i)
        + (n_i - y_i)\log\sigma(-f_i),
    $$

    with site gradient $y_i - n_i\sigma(f_i)$ and Hessian
    $-n_i\sigma(f_i)\sigma(-f_i)$ (log-concave). ``n_trials = 1`` is
    `BernoulliLikelihood`.

    Attributes:
        y: Success counts, shape ``(N,)``.
        n_trials: Numbers of trials $n_i$, a scalar or shape ``(N,)``.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx as gx

        lik = gx.BinomialLikelihood(jnp.array([3.0, 0.0, 5.0]), n_trials=5.0)
        lp = lik.log_prob(jnp.zeros(3))
        grad, hess = lik.site_derivatives(jnp.zeros(3))  # [0.5, -2.5, 2.5], -1.25
        ```
    """

    y: Float[Array, " N"]
    n_trials: Float[Array, " N"] | float

    def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
        """Evaluate the binomial log-likelihood with logit link."""
        y, n = self.y, self.n_trials
        gammaln = jax.scipy.special.gammaln
        log_binom = gammaln(n + 1.0) - gammaln(y + 1.0) - gammaln(n - y + 1.0)
        return jnp.sum(
            log_binom + y * jax.nn.log_sigmoid(f) + (n - y) * jax.nn.log_sigmoid(-f)
        )

    def site_derivatives(
        self, f: Float[Array, " N"]
    ) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
        """Closed-form ``(y − nσ(f), −nσ(f)σ(−f))``."""
        p = jax.nn.sigmoid(f)
        return self.y - self.n_trials * p, -self.n_trials * p * jax.nn.sigmoid(-f)

log_prob(f: Float[Array, ' N']) -> Float[Array, '']

Evaluate the binomial log-likelihood with logit link.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
    """Evaluate the binomial log-likelihood with logit link."""
    y, n = self.y, self.n_trials
    gammaln = jax.scipy.special.gammaln
    log_binom = gammaln(n + 1.0) - gammaln(y + 1.0) - gammaln(n - y + 1.0)
    return jnp.sum(
        log_binom + y * jax.nn.log_sigmoid(f) + (n - y) * jax.nn.log_sigmoid(-f)
    )

site_derivatives(f: Float[Array, ' N']) -> tuple[Float[Array, ' N'], Float[Array, ' N']]

Closed-form (y − nσ(f), −nσ(f)σ(−f)).

Source code in src/gaussx/_quadrature/_likelihoods.py

def site_derivatives(
    self, f: Float[Array, " N"]
) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
    """Closed-form ``(y − nσ(f), −nσ(f)σ(−f))``."""
    p = jax.nn.sigmoid(f)
    return self.y - self.n_trials * p, -self.n_trials * p * jax.nn.sigmoid(-f)

PoissonLikelihood

Bases: AbstractLikelihood

Poisson likelihood with log link.

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Count observations, shape (N,).

Source code in src/gaussx/_quadrature/_likelihoods.py

class PoissonLikelihood(AbstractLikelihood):
    r"""Poisson likelihood with log link.

    Attributes:
        y: Count observations, shape ``(N,)``.
    """

    y: Float[Array, " N"]

    def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
        """Evaluate Poisson log-likelihood with log link."""
        return jnp.sum(self.y * f - jnp.exp(f) - jax.scipy.special.gammaln(self.y + 1))

log_prob(f: Float[Array, ' N']) -> Float[Array, '']

Evaluate Poisson log-likelihood with log link.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
    """Evaluate Poisson log-likelihood with log link."""
    return jnp.sum(self.y * f - jnp.exp(f) - jax.scipy.special.gammaln(self.y + 1))

NegativeBinomialLikelihood

Bases: AbstractLikelihood

Negative-binomial likelihood with log link (NB2, a gamma-Poisson mixture).

\(y_i\) has mean \(\mu_i = e^{f_i}\) and variance \(\mu_i + \mu_i^2/r\), with \(r\) the concentration (R-INLA's size; numpyro's NegativeBinomial2 concentration). \(r\) is the dispersion hyperparameter: a leaf, so it can be traced and differentiated as part of \(\theta\). With \(s_i = \sigma(f_i - \log r) = \mu_i/(r + \mu_i)\),

\[ \log p(y\mid f) = \sum_i \log\frac{\Gamma(y_i + r)}{\Gamma(r)\,y_i!} + r\log r + y_if_i - (r + y_i)\log(r + e^{f_i}), \]

with site gradient \(y_i - (r + y_i)s_i\) and Hessian \(-(r + y_i)s_i(1 - s_i)\) (log-concave). \(r\to\infty\) is PoissonLikelihood.

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Count observations, shape (N,).
concentration	Float[Array, ' N'] | float	\(r > 0\), a scalar or shape (N,).

Examples:

import jax.numpy as jnp
import gaussx as gx

lik = gx.NegativeBinomialLikelihood(jnp.array([0.0, 4.0]), concentration=2.0)
lp = lik.log_prob(jnp.log(jnp.array([1.0, 3.0])))
grad, hess = lik.site_derivatives(jnp.zeros(2))

Source code in src/gaussx/_quadrature/_likelihoods.py

class NegativeBinomialLikelihood(AbstractLikelihood):
    r"""Negative-binomial likelihood with log link (NB2, a gamma-Poisson mixture).

    $y_i$ has mean $\mu_i = e^{f_i}$ and variance $\mu_i + \mu_i^2/r$, with
    $r$ the ``concentration`` (R-INLA's ``size``; numpyro's
    ``NegativeBinomial2`` concentration). $r$ is the dispersion
    hyperparameter: a leaf, so it can be traced and differentiated as part
    of $\theta$. With $s_i = \sigma(f_i - \log r) = \mu_i/(r + \mu_i)$,

    $$
    \log p(y\mid f) = \sum_i \log\frac{\Gamma(y_i + r)}{\Gamma(r)\,y_i!}
        + r\log r + y_if_i - (r + y_i)\log(r + e^{f_i}),
    $$

    with site gradient $y_i - (r + y_i)s_i$ and Hessian
    $-(r + y_i)s_i(1 - s_i)$ (log-concave). $r\to\infty$ is
    `PoissonLikelihood`.

    Attributes:
        y: Count observations, shape ``(N,)``.
        concentration: $r > 0$, a scalar or shape ``(N,)``.

    Examples:
        ```python
        import jax.numpy as jnp
        import gaussx as gx

        lik = gx.NegativeBinomialLikelihood(jnp.array([0.0, 4.0]), concentration=2.0)
        lp = lik.log_prob(jnp.log(jnp.array([1.0, 3.0])))
        grad, hess = lik.site_derivatives(jnp.zeros(2))
        ```
    """

    y: Float[Array, " N"]
    concentration: Float[Array, " N"] | float

    def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
        """Evaluate the negative-binomial log-likelihood with log link."""
        y, r = self.y, self.concentration
        gammaln = jax.scipy.special.gammaln
        log_r = jnp.log(r)
        return jnp.sum(
            gammaln(y + r)
            - gammaln(r)
            - gammaln(y + 1.0)
            + r * log_r
            + y * f
            - (r + y) * jnp.logaddexp(log_r, f)
        )

    def site_derivatives(
        self, f: Float[Array, " N"]
    ) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
        """Closed-form ``(y − (r + y)s, −(r + y)s(1 − s))``, ``s = σ(f − log r)``."""
        y, r = self.y, self.concentration
        shifted = f - jnp.log(r)
        s = jax.nn.sigmoid(shifted)
        return y - (r + y) * s, -(r + y) * s * jax.nn.sigmoid(-shifted)

log_prob(f: Float[Array, ' N']) -> Float[Array, '']

Evaluate the negative-binomial log-likelihood with log link.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
    """Evaluate the negative-binomial log-likelihood with log link."""
    y, r = self.y, self.concentration
    gammaln = jax.scipy.special.gammaln
    log_r = jnp.log(r)
    return jnp.sum(
        gammaln(y + r)
        - gammaln(r)
        - gammaln(y + 1.0)
        + r * log_r
        + y * f
        - (r + y) * jnp.logaddexp(log_r, f)
    )

site_derivatives(f: Float[Array, ' N']) -> tuple[Float[Array, ' N'], Float[Array, ' N']]

Closed-form (y − (r + y)s, −(r + y)s(1 − s)), s = σ(f − log r).

Source code in src/gaussx/_quadrature/_likelihoods.py

def site_derivatives(
    self, f: Float[Array, " N"]
) -> tuple[Float[Array, " N"], Float[Array, " N"]]:
    """Closed-form ``(y − (r + y)s, −(r + y)s(1 − s))``, ``s = σ(f − log r)``."""
    y, r = self.y, self.concentration
    shifted = f - jnp.log(r)
    s = jax.nn.sigmoid(shifted)
    return y - (r + y) * s, -(r + y) * s * jax.nn.sigmoid(-shifted)

SoftmaxLikelihood

Bases: AbstractLikelihood

Softmax (categorical) likelihood for multi-class classification.

Parameters:

Name	Type	Description	Default
y	Int[Array, ' N']	Integer class labels, shape (N,).	required
num_classes	int	Number of classes C.	required

Source code in src/gaussx/_quadrature/_likelihoods.py

class SoftmaxLikelihood(AbstractLikelihood):
    r"""Softmax (categorical) likelihood for multi-class classification.

    Args:
        y: Integer class labels, shape ``(N,)``.
        num_classes: Number of classes C.
    """

    y: Int[Array, " N"]
    num_classes: int = eqx.field(static=True)
    latent_dim: int = eqx.field(static=True, default=1)

    def __init__(self, y: Int[Array, " N"], num_classes: int):
        self.y = y
        self.num_classes = num_classes
        self.latent_dim = num_classes

    def log_prob(self, f: Float[Array, " NC"]) -> Float[Array, ""]:
        """Evaluate softmax log-likelihood."""
        f_2d = rearrange(f, "(N C) -> N C", C=self.num_classes)
        log_probs = jax.nn.log_softmax(f_2d, axis=-1)
        return jnp.sum(log_probs[jnp.arange(self.y.shape[0]), self.y])

log_prob(f: Float[Array, ' NC']) -> Float[Array, '']

Evaluate softmax log-likelihood.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " NC"]) -> Float[Array, ""]:
    """Evaluate softmax log-likelihood."""
    f_2d = rearrange(f, "(N C) -> N C", C=self.num_classes)
    log_probs = jax.nn.log_softmax(f_2d, axis=-1)
    return jnp.sum(log_probs[jnp.arange(self.y.shape[0]), self.y])

StudentTLikelihood

Bases: AbstractLikelihood

Student-t likelihood for robust regression.

Attributes:

Name	Type	Description
y	Float[Array, ' N']	Observations, shape (N,).
df	float	Degrees of freedom (> 0).
scale	float	Scale parameter (> 0).

Source code in src/gaussx/_quadrature/_likelihoods.py

class StudentTLikelihood(AbstractLikelihood):
    r"""Student-t likelihood for robust regression.

    Attributes:
        y: Observations, shape ``(N,)``.
        df: Degrees of freedom (> 0).
        scale: Scale parameter (> 0).
    """

    y: Float[Array, " N"]
    df: float
    scale: float

    def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
        """Evaluate Student-t log-likelihood."""
        df = self.df
        scale = self.scale
        residual = self.y - f
        half_df = 0.5 * df
        half_dfp1 = 0.5 * (df + 1.0)

        log_norm = (
            jax.scipy.special.gammaln(half_dfp1)
            - jax.scipy.special.gammaln(half_df)
            - 0.5 * jnp.log(df * jnp.pi * scale**2)
        )
        log_kernel = -half_dfp1 * jnp.log1p(residual**2 / (df * scale**2))
        return jnp.sum(log_norm + log_kernel)

log_prob(f: Float[Array, ' N']) -> Float[Array, '']

Evaluate Student-t log-likelihood.

Source code in src/gaussx/_quadrature/_likelihoods.py

def log_prob(self, f: Float[Array, " N"]) -> Float[Array, ""]:
    """Evaluate Student-t log-likelihood."""
    df = self.df
    scale = self.scale
    residual = self.y - f
    half_df = 0.5 * df
    half_dfp1 = 0.5 * (df + 1.0)

    log_norm = (
        jax.scipy.special.gammaln(half_dfp1)
        - jax.scipy.special.gammaln(half_df)
        - 0.5 * jnp.log(df * jnp.pi * scale**2)
    )
    log_kernel = -half_dfp1 * jnp.log1p(residual**2 / (df * scale**2))
    return jnp.sum(log_norm + log_kernel)

Expectations & EP moments

Expected log-likelihoods (the ELL term of every ELBO), generic mean / gradient / cost expectations, and the tilted-moment matching at the heart of expectation propagation.

Structured linear algebra and Gaussian primitives for JAX.

elbo(likelihood: AbstractLikelihood, state: GaussianState, kl: Float[Array, ''], integrator: AbstractIntegrator | None = None) -> Float[Array, '']

Evidence lower bound (ELBO).

Computes:

ELBO = E_q[log p(y | f)] - KL(q || p)

Dispatches to analytical expected log-likelihood when available (e.g. Gaussian likelihood), or uses numerical integration.

Parameters:

Name	Type	Description	Default
likelihood	AbstractLikelihood	Likelihood object with log_prob method.	required
state	GaussianState	Variational Gaussian distribution q(f).	required
kl	Float[Array, '']	KL divergence KL(q || p) (scalar, precomputed).	required
integrator	AbstractIntegrator | None	Integration method for non-conjugate likelihoods. Ignored when the likelihood has an analytical fast path.	None

Returns:

Type	Description
Float[Array, '']	Scalar ELBO value.

Source code in src/gaussx/_quadrature/_expectations.py

def elbo(
    likelihood: AbstractLikelihood,
    state: GaussianState,
    kl: Float[Array, ""],
    integrator: AbstractIntegrator | None = None,
) -> Float[Array, ""]:
    r"""Evidence lower bound (ELBO).

    Computes:

        ELBO = E_q[log p(y | f)] - KL(q || p)

    Dispatches to analytical expected log-likelihood when available
    (e.g. Gaussian likelihood), or uses numerical integration.

    Args:
        likelihood: Likelihood object with ``log_prob`` method.
        state: Variational Gaussian distribution ``q(f)``.
        kl: KL divergence ``KL(q || p)`` (scalar, precomputed).
        integrator: Integration method for non-conjugate likelihoods.
            Ignored when the likelihood has an analytical fast path.

    Returns:
        Scalar ELBO value.
    """
    ell = expected_log_likelihood(likelihood, state, integrator)
    return ell - kl

expected_log_likelihood(likelihood: AbstractLikelihood, state: GaussianState, integrator: AbstractIntegrator | None = None) -> Float[Array, '']

Unified expected log-likelihood with analytical dispatch.

Computes E_q[log p(y | f)] where q = N(mu, Sigma).

If the likelihood has a closed-form expected log-likelihood (e.g. GaussianLikelihood), it is used directly without an integrator. Otherwise, an integrator must be provided for numerical approximation.

Parameters:

Name	Type	Description	Default
likelihood	AbstractLikelihood	Likelihood object with log_prob method.	required
state	GaussianState	Variational Gaussian distribution.	required
integrator	AbstractIntegrator | None	Integration method. Required for non-conjugate likelihoods; ignored when the likelihood has an analytical fast path.	None

Returns:

Type	Description
Float[Array, '']	Scalar expected log-likelihood.

Raises:

Type	Description
ValueError	If no integrator is provided and the likelihood has no analytical form.

Source code in src/gaussx/_quadrature/_expectations.py

def expected_log_likelihood(
    likelihood: AbstractLikelihood,
    state: GaussianState,
    integrator: AbstractIntegrator | None = None,
) -> Float[Array, ""]:
    r"""Unified expected log-likelihood with analytical dispatch.

    Computes ``E_q[log p(y | f)]`` where ``q = N(mu, Sigma)``.

    If the likelihood has a closed-form expected log-likelihood
    (e.g. ``GaussianLikelihood``), it is used directly without
    an integrator. Otherwise, an ``integrator`` must be provided
    for numerical approximation.

    Args:
        likelihood: Likelihood object with ``log_prob`` method.
        state: Variational Gaussian distribution.
        integrator: Integration method. Required for non-conjugate
            likelihoods; ignored when the likelihood has an
            analytical fast path.

    Returns:
        Scalar expected log-likelihood.

    Raises:
        ValueError: If no integrator is provided and the likelihood
            has no analytical form.
    """
    if likelihood.has_analytical_ell():
        return likelihood.analytical_expected_log_likelihood(state.mean, state.cov)
    if integrator is None:
        msg = (
            f"{type(likelihood).__name__} has no analytical ELL; provide an integrator"
        )
        raise ValueError(msg)
    return log_likelihood_expectation(likelihood.log_prob, state, integrator)

log_likelihood_expectation(likelihood_fn: Callable[[Float[Array, ' N']], Float[Array, '']], state: GaussianState, integrator: AbstractIntegrator) -> Float[Array, '']

Compute E[log p(y_{obs} | f(x))] where x ~ N(mu, Sigma).

For non-conjugate likelihoods (Bernoulli, Poisson, etc.) where the expectation has no closed form.

Parameters:

Name	Type	Description	Default
likelihood_fn	Callable[[Float[Array, ' N']], Float[Array, '']]	Function mapping latent values to scalar log-likelihood: (N,) -> ().	required
state	GaussianState	Input Gaussian distribution.	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
Float[Array, '']	Scalar expected log-likelihood.

Source code in src/gaussx/_quadrature/_expectations.py

def log_likelihood_expectation(
    likelihood_fn: Callable[[Float[Array, " N"]], Float[Array, ""]],
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> Float[Array, ""]:
    r"""Compute ``E[log p(y_{obs} | f(x))]`` where ``x ~ N(mu, Sigma)``.

    For non-conjugate likelihoods (Bernoulli, Poisson, etc.) where
    the expectation has no closed form.

    Args:
        likelihood_fn: Function mapping latent values to scalar
            log-likelihood: ``(N,) -> ()``.
        state: Input Gaussian distribution.
        integrator: Integration method.

    Returns:
        Scalar expected log-likelihood.
    """
    return mean_expectation(
        lambda x: jnp.atleast_1d(likelihood_fn(x)),
        state,
        integrator,
    )[0]

mean_expectation(fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], state: GaussianState, integrator: AbstractIntegrator) -> Float[Array, ' M']

Compute E[f(x)] where x ~ N(mu, Sigma).

For a point-based rule (one that implements points_and_weights) this is the weighted sum \sum_p w^{(m)}_p f(\chi_p) alone: the output covariance that integrator.integrate would also assemble is never built, which matters for wide outputs (M = N^2 for a flattened outer product). Other integrators (e.g. Taylor) fall back to integrator.integrate(fn, state).state.mean.

Parameters:

Name	Type	Description	Default
fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	Function mapping (N,) -> (M,).	required
state	GaussianState	Input Gaussian distribution.	required
integrator	AbstractIntegrator	Integration method (Taylor, Unscented, MC, etc.).	required

Returns:

Type	Description
Float[Array, ' M']	Expected function value, shape (M,).

Source code in src/gaussx/_quadrature/_expectations.py

def mean_expectation(
    fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> Float[Array, " M"]:
    r"""Compute ``E[f(x)]`` where ``x ~ N(mu, Sigma)``.

    For a point-based rule (one that implements ``points_and_weights``)
    this is the weighted sum ``\sum_p w^{(m)}_p f(\chi_p)`` alone: the
    output covariance that ``integrator.integrate`` would also assemble is
    never built, which matters for wide outputs (``M = N^2`` for a
    flattened outer product). Other integrators (e.g. Taylor) fall back to
    ``integrator.integrate(fn, state).state.mean``.

    Args:
        fn: Function mapping ``(N,) -> (M,)``.
        state: Input Gaussian distribution.
        integrator: Integration method (Taylor, Unscented, MC, etc.).

    Returns:
        Expected function value, shape ``(M,)``.
    """
    rule = _points_and_weights_or_none(integrator, state)
    if rule is None:
        return integrator.integrate(fn, state).state.mean
    chi, w_m, _ = rule
    return einsum(w_m, jax.vmap(fn)(chi), "p, p m -> m")

gradient_expectation(fn: Callable[[Float[Array, ' N']], Float[Array, '']], state: GaussianState, integrator: AbstractIntegrator) -> Float[Array, ' N']

Compute E[nabla f(x)] via Stein's lemma.

Uses the identity:

E[nabla f(x)] = Sigma^{-1} Cov[x, f(x)]

Parameters:

Name	Type	Description	Default
fn	Callable[[Float[Array, ' N']], Float[Array, '']]	Scalar-valued function mapping (N,) -> ().	required
state	GaussianState	Input Gaussian distribution.	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
Float[Array, ' N']	Expected gradient, shape (N,).

Source code in src/gaussx/_quadrature/_expectations.py

def gradient_expectation(
    fn: Callable[[Float[Array, " N"]], Float[Array, ""]],
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> Float[Array, " N"]:
    r"""Compute ``E[nabla f(x)]`` via Stein's lemma.

    Uses the identity:

        E[nabla f(x)] = Sigma^{-1} Cov[x, f(x)]

    Args:
        fn: Scalar-valued function mapping ``(N,) -> ()``.
        state: Input Gaussian distribution.
        integrator: Integration method.

    Returns:
        Expected gradient, shape ``(N,)``.
    """

    # Wrap scalar fn to return (1,) for the integrator
    def fn_vec(x: Float[Array, " N"]) -> Float[Array, " 1"]:
        return jnp.atleast_1d(fn(x))

    result = integrator.integrate(fn_vec, state)
    cross_cov = result.cross_cov
    assert cross_cov is not None, "Integrator must return cross_cov"
    # E[nabla f] = Sigma^{-1} @ Cov[x, f]  (N, 1) -> (N,)
    return solve(state.cov, cross_cov[:, 0])

cost_expectation(prediction_fn: Callable[[Float[Array, ' N']], Float[Array, ' M']], cost_fn: Callable[[Float[Array, ' M'], Float[Array, ' M']], Float[Array, '']], state: GaussianState, target: Float[Array, ' M'], integrator: AbstractIntegrator) -> Float[Array, '']

Compute E[Cost(f(x), target)] where x ~ N(mu, Sigma).

For model-based RL: expected cost of a policy under state uncertainty.

Parameters:

Name	Type	Description	Default
prediction_fn	Callable[[Float[Array, ' N']], Float[Array, ' M']]	Maps state to prediction, (N,) -> (M,).	required
cost_fn	Callable[[Float[Array, ' M'], Float[Array, ' M']], Float[Array, '']]	Cost function, (M,), (M,) -> ().	required
state	GaussianState	Input Gaussian distribution (uncertain state).	required
target	Float[Array, ' M']	Target value, shape (M,).	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
Float[Array, '']	Scalar expected cost.

Source code in src/gaussx/_quadrature/_expectations.py

def cost_expectation(
    prediction_fn: Callable[[Float[Array, " N"]], Float[Array, " M"]],
    cost_fn: Callable[[Float[Array, " M"], Float[Array, " M"]], Float[Array, ""]],
    state: GaussianState,
    target: Float[Array, " M"],
    integrator: AbstractIntegrator,
) -> Float[Array, ""]:
    r"""Compute ``E[Cost(f(x), target)]`` where ``x ~ N(mu, Sigma)``.

    For model-based RL: expected cost of a policy under state
    uncertainty.

    Args:
        prediction_fn: Maps state to prediction, ``(N,) -> (M,)``.
        cost_fn: Cost function, ``(M,), (M,) -> ()``.
        state: Input Gaussian distribution (uncertain state).
        target: Target value, shape ``(M,)``.
        integrator: Integration method.

    Returns:
        Scalar expected cost.
    """

    def combined_fn(x: Float[Array, " N"]) -> Float[Array, " 1"]:
        pred = prediction_fn(x)
        return jnp.atleast_1d(cost_fn(pred, target))

    return mean_expectation(combined_fn, state, integrator)[0]

ep_tilted_moments(log_lik_fn: Callable[[Float[Array, '']], Float[Array, '']], cav_mean: Float[Array, ' *batch'], cav_var: Float[Array, ' *batch'], *, order: int | None = None, integrator: AbstractIntegrator | None = None, power: float = 1.0) -> tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]

Moments of the scalar EP tilted distribution, elementwise over a batch.

For each cavity N(f | m, v) computes the mean and variance of the tilted distribution

\[ \tilde p(f) \propto p(y \mid f)^{a}\, \mathcal{N}(f \mid m, v) \]

with the points f_i and mean weights w_i of a point-based rule applied to the cavity:

\[ \tilde w_i = w_i\, p(y \mid f_i)^{a}, \quad m_{\mathrm{tilt}} = \frac{\sum_i \tilde w_i f_i}{\sum_i \tilde w_i}, \quad v_{\mathrm{tilt}} = \frac{\sum_i \tilde w_i (f_i - m_{\mathrm{tilt}})^2} {\sum_i \tilde w_i}. \]

These are the same tilted moments as m + v g and v + v H v from gaussx.moment_match (Stein's lemma), and share its log-sum-exp shifted weights. The default rule is Gauss-Hermite of order points, the behaviour before integrator existed.

Pseudocode (per batch element):

if v <= 0: return m, v                      # invalid site, passed through
f, w = integrator.points_and_weights(N(m, v))
w~ = w * exp(a * log_lik(f) - max(a * log_lik(f)))
m_t = sum(w~ f) / sum(w~);  v_t = sum(w~ (f - m_t)^2) / sum(w~)
return m_t, max(v_t, eps * v)

A site whose cavity variance is not positive -- EP's usual failure mode, when the site holds more precision than the current posterior -- has no tilted distribution. Its cavity moments are returned unchanged, so the moment-matched site update 1/t_var - 1/cav_var is zero and nothing NaN reaches the site parameters; skip or damp such sites in the caller. Valid sites are unaffected, and gradients stay finite for both.

Parameters:

Name	Type	Description	Default
log_lik_fn	Callable[[Float[Array, '']], Float[Array, '']]	Scalar function mapping latent value f to scalar log-likelihood log p(y|f).	required
cav_mean	Float[Array, ' *batch']	Cavity means, shape (*batch,).	required
cav_var	Float[Array, ' *batch']	Cavity variances, shape (*batch,). Non-positive entries are passed through as described above.	required
order	int | None	Number of Gauss-Hermite points of the default rule. Default 20. Only valid when integrator is None.	None
integrator	AbstractIntegrator | None	Point-based rule (e.g. gaussx.UnscentedIntegrator, gaussx.CubatureIntegrator) applied to each 1-D cavity. Default GaussHermiteIntegrator(order).	None
power	float	Power-EP exponent a; 1.0 is standard EP.	1.0

Returns:

Type	Description
tuple[Float[Array, ' *batch'], Float[Array, ' *batch']]	Tuple (tilted_mean, tilted_var), each of shape (*batch,).

Raises:

Type	Description
ValueError	If both integrator and order are given.
NotImplementedError	If integrator is not point-based.

Source code in src/gaussx/_quadrature/_tilted_moments.py

def ep_tilted_moments(
    log_lik_fn: Callable[[Float[Array, ""]], Float[Array, ""]],
    cav_mean: Float[Array, " *batch"],
    cav_var: Float[Array, " *batch"],
    *,
    order: int | None = None,
    integrator: AbstractIntegrator | None = None,
    power: float = 1.0,
) -> tuple[Float[Array, " *batch"], Float[Array, " *batch"]]:
    r"""Moments of the scalar EP tilted distribution, elementwise over a batch.

    For each cavity ``N(f | m, v)`` computes the mean and variance of the
    tilted distribution

    $$
    \tilde p(f) \propto p(y \mid f)^{a}\, \mathcal{N}(f \mid m, v)
    $$

    with the points ``f_i`` and mean weights ``w_i`` of a point-based rule
    applied to the cavity:

    $$
    \tilde w_i = w_i\, p(y \mid f_i)^{a}, \quad
    m_{\mathrm{tilt}} = \frac{\sum_i \tilde w_i f_i}{\sum_i \tilde w_i}, \quad
    v_{\mathrm{tilt}} = \frac{\sum_i \tilde w_i (f_i - m_{\mathrm{tilt}})^2}
        {\sum_i \tilde w_i}.
    $$

    These are the same tilted moments as ``m + v g`` and ``v + v H v`` from
    `gaussx.moment_match` (Stein's lemma), and share its log-sum-exp shifted
    weights. The default rule is Gauss-Hermite of ``order`` points, the
    behaviour before ``integrator`` existed.

    Pseudocode (per batch element):

    ```text
    if v <= 0: return m, v                      # invalid site, passed through
    f, w = integrator.points_and_weights(N(m, v))
    w~ = w * exp(a * log_lik(f) - max(a * log_lik(f)))
    m_t = sum(w~ f) / sum(w~);  v_t = sum(w~ (f - m_t)^2) / sum(w~)
    return m_t, max(v_t, eps * v)
    ```

    A site whose cavity variance is not positive -- EP's usual failure mode,
    when the site holds more precision than the current posterior -- has no
    tilted distribution. Its cavity moments are returned unchanged, so the
    moment-matched site update ``1/t_var - 1/cav_var`` is zero and nothing
    ``NaN`` reaches the site parameters; skip or damp such sites in the
    caller. Valid sites are unaffected, and gradients stay finite for both.

    Args:
        log_lik_fn: Scalar function mapping latent value ``f`` to scalar
            log-likelihood ``log p(y|f)``.
        cav_mean: Cavity means, shape ``(*batch,)``.
        cav_var: Cavity variances, shape ``(*batch,)``. Non-positive entries
            are passed through as described above.
        order: Number of Gauss-Hermite points of the default rule. Default
            ``20``. Only valid when ``integrator`` is ``None``.
        integrator: Point-based rule (e.g. `gaussx.UnscentedIntegrator`,
            `gaussx.CubatureIntegrator`) applied to each 1-D cavity. Default
            ``GaussHermiteIntegrator(order)``.
        power: Power-EP exponent ``a``; ``1.0`` is standard EP.

    Returns:
        Tuple ``(tilted_mean, tilted_var)``, each of shape ``(*batch,)``.

    Raises:
        ValueError: If both ``integrator`` and ``order`` are given.
        NotImplementedError: If ``integrator`` is not point-based.
    """
    if integrator is None:
        integrator = GaussHermiteIntegrator(order=20 if order is None else order)
    elif order is not None:
        msg = "Pass either `integrator` or `order` (Gauss-Hermite), not both."
        raise ValueError(msg)

    # Cavity (and so the rule's points) in at least float32, without
    # promoting a float32 cavity to float64 under x64. promote_types (not
    # result_type with a float32 operand) so a weakly typed float64 cavity,
    # e.g. jnp.array(0.3), stays float64.
    dtype = jnp.promote_types(jnp.result_type(cav_mean, cav_var), jnp.float32)
    cav_mean = jnp.asarray(cav_mean, dtype=dtype)
    cav_var = jnp.asarray(cav_var, dtype=dtype)

    def _quadrature_moments(mean_i: Float[Array, ""], var_i: Float[Array, ""]):
        state = GaussianState(
            mean=jnp.atleast_1d(mean_i),
            cov=lx.DiagonalLinearOperator(jnp.atleast_1d(var_i)),
        )
        chi, w_m, _ = integrator.points_and_weights(state)  # (P, 1), (P,)
        p_tilde, _ = _tilted_weights(lambda x: log_lik_fn(x[0]), chi, w_m, power)
        f_nodes = chi[:, 0]
        Z = jnp.sum(p_tilde)
        t_mean = jnp.sum(p_tilde * f_nodes) / Z
        t_var = jnp.sum(p_tilde * (f_nodes - t_mean) ** 2) / Z
        # Floor relative to the cavity, not in absolute units: a tilted
        # variance below eps * cavity variance is quadrature round-off, and
        # an absolute floor would inflate small-scale problems.
        t_var = jnp.maximum(t_var, jnp.finfo(t_var.dtype).eps * var_i)
        return t_mean, t_var

    def _compute_moments(mean_i: Float[Array, ""], var_i: Float[Array, ""]):
        # Double where: the quadrature never sees a negative variance, so
        # neither the value nor the gradient of an invalid site is NaN.
        valid = var_i > 0
        t_mean, t_var = _quadrature_moments(mean_i, jnp.where(valid, var_i, 1.0))
        return jnp.where(valid, t_mean, mean_i), jnp.where(valid, t_var, var_i)

    return jnp.vectorize(_compute_moments)(cav_mean, cav_var)

Moment matching & linearisation

moment_match returns the tilted log-normaliser and its cavity-mean derivatives from a single pass over the quadrature points; the tilted moments and the EP site follow from those in closed form (see its docstring — the derivatives are not themselves the site naturals). statistical_linear_regression returns the moment-matched linear-Gaussian surrogate \(p(y \mid f) \approx \mathcal{N}(y \mid Af + b, \Omega)\) behind posterior linearisation and the iterated Kalman smoother. Both take any point-based integrator.

Structured linear algebra and Gaussian primitives for JAX.

MomentMatchResult

Bases: Module

Tilted log-normaliser and its derivatives w.r.t. the cavity mean.

Attributes:

Name	Type	Description
log_Z	Float[Array, '']	Scalar tilted log-normaliser log Z.
d_log_Z	Float[Array, ' D']	Gradient ∂ log Z / ∂m, shape (D,).
d2_log_Z	Float[Array, 'D D']	Hessian ∂² log Z / ∂m², shape (D, D).

Source code in src/gaussx/_quadrature/_moment_match.py

class MomentMatchResult(eqx.Module):
    r"""Tilted log-normaliser and its derivatives w.r.t. the cavity mean.

    Attributes:
        log_Z: Scalar tilted log-normaliser ``log Z``.
        d_log_Z: Gradient ``∂ log Z / ∂m``, shape ``(D,)``.
        d2_log_Z: Hessian ``∂² log Z / ∂m²``, shape ``(D, D)``.
    """

    log_Z: Float[Array, ""]
    d_log_Z: Float[Array, " D"]
    d2_log_Z: Float[Array, "D D"]

SLRResult

Bases: Module

Linear-Gaussian surrogate produced by statistical linear regression.

Encodes p(y | f) ≈ N(y | A f + b, Ω), the moment-matched approximation of a nonlinear conditional under a Gaussian state.

Attributes:

Name	Type	Description
A	Float[Array, 'M D']	Linearisation gain, shape (M, D).
b	Float[Array, ' M']	Linearisation offset, shape (M,).
omega	Float[Array, 'M M']	Residual noise covariance, shape (M, M).
mu	Float[Array, ' M']	Expected conditional mean E_q[E[y | f]], shape (M,).

Source code in src/gaussx/_quadrature/_slr.py

class SLRResult(eqx.Module):
    r"""Linear-Gaussian surrogate produced by statistical linear regression.

    Encodes ``p(y | f) ≈ N(y | A f + b, Ω)``, the moment-matched
    approximation of a nonlinear conditional under a Gaussian state.

    Attributes:
        A: Linearisation gain, shape ``(M, D)``.
        b: Linearisation offset, shape ``(M,)``.
        omega: Residual noise covariance, shape ``(M, M)``.
        mu: Expected conditional mean ``E_q[E[y | f]]``, shape ``(M,)``.
    """

    A: Float[Array, "M D"]
    b: Float[Array, " M"]
    omega: Float[Array, "M M"]
    mu: Float[Array, " M"]

moment_match(log_lik_fn: Callable[[Float[Array, ' D']], Float[Array, '']], state: GaussianState, integrator: AbstractIntegrator, *, power: float = 1.0) -> MomentMatchResult

EP moment matching: tilted log-normaliser and its first two derivatives.

Computes the log normaliser of the tilted distribution,

\[ \log Z = \log \int p(y \mid f)^{a}\, \mathcal{N}(f \mid m, C)\, df, \]

together with

\[ \partial_m \log Z = C^{-1} \mathbb{E}_{\mathrm{tilt}}[f - m], \qquad \partial^2_m \log Z = C^{-1}\big(\mathrm{Cov}_{\mathrm{tilt}}[f] - C\big) C^{-1}. \]

This is the EP quantity, not E_q[log p(y|f)] (which is what variational inference needs — see gaussx.expected_log_likelihood).

Writing g = d_log_Z and H = d2_log_Z, the tilted moments follow directly from the two derivatives:

\[ m_{\mathrm{tilt}} = m + C g, \qquad C_{\mathrm{tilt}} = C + C H C. \]

The Gaussian site is the natural-parameter difference between the tilted distribution and the cavity, which reduces to

\[ \Lambda_{\mathrm{site}} = C_{\mathrm{tilt}}^{-1} - C^{-1} = -(I + H C)^{-1} H, \qquad \lambda_{\mathrm{site}} = C_{\mathrm{tilt}}^{-1} m_{\mathrm{tilt}} - C^{-1} m = (I + H C)^{-1} (g - H m), \]

in the same nat2 = +Λ convention gaussx.cavity_distribution uses. That site approximates p(y | f) ** power, so divide both naturals by power to recover the site for p(y | f) itself.

Do not feed (g, -H) in as the site: it drops the -H m term and the (I + H C)^{-1} factor. Nor is gaussx.newton_update the right conversion here — it yields (g - H m, -H), which is the local-Newton site, i.e. only the first-order approximation to the EP site above.

All three outputs come from a single pass over the quadrature points: applying Stein's lemma to the derivatives of the Gaussian density turns them into reweightings of the same likelihood evaluations, so no autodiff runs through the integrator and no extra likelihood evaluations are needed.

The sum over points is evaluated with the log-sum-exp shift, so likelihoods that underflow in the raw scale are handled. Note that quadrature weights may be negative (the scaled unscented transform and the degree-5 cubature rule both have negative weights in places), so Z is not guaranteed positive for a badly matched rule; prefer a rule whose points cover the likelihood's mass.

Parameters:

Name	Type	Description	Default
log_lik_fn	Callable[[Float[Array, ' D']], Float[Array, '']]	Maps a latent f of shape (D,) to the scalar log-likelihood log p(y | f).	required
state	GaussianState	Cavity distribution N(m, C). Its covariance must be symmetric.	required
integrator	AbstractIntegrator	Point-based rule supplying the points and weights, e.g. gaussx.FifthOrderCubatureIntegrator.	required
power	float	Power-EP exponent a. 1.0 is standard EP; a < 1 corresponds to fractional EP / alpha-divergence minimisation. Documented for a in (0, 1].	1.0

Returns:

Type	Description
MomentMatchResult	MomentMatchResult with log_Z, d_log_Z, d2_log_Z.

Raises:

Type	Description
NotImplementedError	If integrator is not point-based.
ValueError	If log_lik_fn does not return a scalar per point.

Source code in src/gaussx/_quadrature/_moment_match.py

def moment_match(
    log_lik_fn: Callable[[Float[Array, " D"]], Float[Array, ""]],
    state: GaussianState,
    integrator: AbstractIntegrator,
    *,
    power: float = 1.0,
) -> MomentMatchResult:
    r"""EP moment matching: tilted log-normaliser and its first two derivatives.

    Computes the log normaliser of the tilted distribution,

    $$
    \log Z = \log \int p(y \mid f)^{a}\, \mathcal{N}(f \mid m, C)\, df,
    $$

    together with

    $$
    \partial_m \log Z = C^{-1} \mathbb{E}_{\mathrm{tilt}}[f - m],
    \qquad
    \partial^2_m \log Z
        = C^{-1}\big(\mathrm{Cov}_{\mathrm{tilt}}[f] - C\big) C^{-1}.
    $$

    This is the EP quantity, *not* ``E_q[log p(y|f)]`` (which is what
    variational inference needs — see `gaussx.expected_log_likelihood`).

    Writing ``g = d_log_Z`` and ``H = d2_log_Z``, the tilted moments follow
    directly from the two derivatives:

    $$
    m_{\mathrm{tilt}} = m + C g, \qquad
    C_{\mathrm{tilt}} = C + C H C.
    $$

    The Gaussian site is the natural-parameter *difference* between the
    tilted distribution and the cavity, which reduces to

    $$
    \Lambda_{\mathrm{site}} = C_{\mathrm{tilt}}^{-1} - C^{-1}
        = -(I + H C)^{-1} H,
    \qquad
    \lambda_{\mathrm{site}} = C_{\mathrm{tilt}}^{-1} m_{\mathrm{tilt}}
        - C^{-1} m = (I + H C)^{-1} (g - H m),
    $$

    in the same ``nat2 = +Λ`` convention `gaussx.cavity_distribution` uses.
    That site approximates ``p(y | f) ** power``, so divide both naturals by
    ``power`` to recover the site for ``p(y | f)`` itself.

    Do **not** feed ``(g, -H)`` in as the site: it drops the ``-H m`` term
    and the ``(I + H C)^{-1}`` factor. Nor is `gaussx.newton_update` the
    right conversion here — it yields ``(g - H m, -H)``, which is the
    local-Newton site, i.e. only the first-order approximation to the EP
    site above.

    All three outputs come from a **single** pass over the quadrature
    points: applying Stein's lemma to the derivatives of the Gaussian
    density turns them into reweightings of the same likelihood
    evaluations, so no autodiff runs through the integrator and no extra
    likelihood evaluations are needed.

    The sum over points is evaluated with the log-sum-exp shift, so
    likelihoods that underflow in the raw scale are handled. Note that
    quadrature weights may be negative (the scaled unscented transform and
    the degree-5 cubature rule both have negative weights in places), so
    ``Z`` is not guaranteed positive for a badly matched rule; prefer a
    rule whose points cover the likelihood's mass.

    Args:
        log_lik_fn: Maps a latent ``f`` of shape ``(D,)`` to the scalar
            log-likelihood ``log p(y | f)``.
        state: Cavity distribution ``N(m, C)``. Its covariance must be
            symmetric.
        integrator: Point-based rule supplying the points and weights,
            e.g. `gaussx.FifthOrderCubatureIntegrator`.
        power: Power-EP exponent ``a``. ``1.0`` is standard EP; ``a < 1``
            corresponds to fractional EP / alpha-divergence minimisation.
            Documented for ``a`` in ``(0, 1]``.

    Returns:
        `MomentMatchResult` with ``log_Z``, ``d_log_Z``, ``d2_log_Z``.

    Raises:
        NotImplementedError: If ``integrator`` is not point-based.
        ValueError: If ``log_lik_fn`` does not return a scalar per point.
    """
    chi, w_m, _ = integrator.points_and_weights(state)
    p_tilde, shift = _tilted_weights(log_lik_fn, chi, w_m, power)  # (P,), signed
    Z_tilde = jnp.sum(p_tilde)

    log_Z = shift + jnp.log(Z_tilde)

    # u_i = C^{-1} (f_i − m); C is symmetric so a row-wise solve suffices.
    dx = chi - state.mean[None, :]  # (P, D)
    U = solve_rows(state.cov, dx)  # (P, D)

    # ∂_m log Z = Z^{-1} Σ w_i u_i p_i
    d_log_Z = jnp.sum(p_tilde[:, None] * U, axis=0) / Z_tilde

    # Z^{-1} ∂²_m Z = Z^{-1} Σ w_i (u_i u_iᵀ − C^{-1}) p_i
    second = (
        jnp.sum(
            p_tilde[:, None, None] * (U[:, :, None] * U[:, None, :]),
            axis=0,
        )
        / Z_tilde
    )
    cov_inv = inv(state.cov).as_matrix()

    # ∂²_m log Z = Z^{-1} ∂²_m Z − (∂_m log Z)(∂_m log Z)ᵀ
    d2_log_Z = symmetrize(second - cov_inv - jnp.outer(d_log_Z, d_log_Z))

    return MomentMatchResult(log_Z=log_Z, d_log_Z=d_log_Z, d2_log_Z=d2_log_Z)

statistical_linear_regression(conditional_mean_fn: Callable[[Float[Array, ' D']], Float[Array, ' M']], conditional_var_fn: Callable[[Float[Array, ' D']], Float[Array, 'M M']], state: GaussianState, integrator: AbstractIntegrator) -> SLRResult

Moment-match a nonlinear conditional to a linear-Gaussian surrogate.

Given a conditional p(y | f) described by its first two moments and a Gaussian state q(f) = N(m_q, C_ff), returns the linear-Gaussian model p(y | f) ≈ N(y | A f + b, Ω) whose moments match those of the true conditional under q.

With

\[ \mu = \mathbb{E}_q[\mathbb{E}[y \mid f]], \quad C_{xf} = \mathrm{Cov}_q(f, \mathbb{E}[y \mid f]), \quad S = \mathrm{Var}_q(\mathbb{E}[y \mid f]) + \mathbb{E}_q[\mathrm{Var}[y \mid f]], \]

the surrogate is

\[ A = C_{xf}^\top C_{ff}^{-1}, \quad b = \mu - A m_q, \quad \Omega = S - A C_{xf}. \]

This is the shared primitive behind posterior linearisation, the iterated Kalman smoother (IPLF / IPLS), and the conjugate Gaussian site that an EP cavity meets. Feeding (A, b, omega) into a standard linear-Gaussian Kalman update reproduces posterior-linearisation inference exactly.

All expectations are evaluated with integrator, so the accuracy / cost trade-off is chosen by the caller. Only the conditional moments are needed — no autodiff runs through the quadrature.

Parameters:

Name	Type	Description	Default
conditional_mean_fn	Callable[[Float[Array, ' D']], Float[Array, ' M']]	Maps a latent f of shape (D,) to the conditional mean E[y | f] of shape (M,).	required
conditional_var_fn	Callable[[Float[Array, ' D']], Float[Array, 'M M']]	Maps a latent f of shape (D,) to the conditional covariance Var[y | f]. Either the full (M, M) covariance or, for conditionally independent outputs, an (M,) vector of variances.	required
state	GaussianState	Gaussian state q(f). Its covariance must be symmetric.	required
integrator	AbstractIntegrator	Point-based rule used for the expectations, e.g. gaussx.FifthOrderCubatureIntegrator.	required

Returns:

Type	Description
SLRResult	SLRResult with the gain, offset, residual covariance, and
SLRResult	expected conditional mean.

Raises:

Type	Description
NotImplementedError	If integrator is not point-based.
ValueError	If conditional_var_fn does not return an (M,) or (M, M) array matching the M outputs of conditional_mean_fn.

Source code in src/gaussx/_quadrature/_slr.py

def statistical_linear_regression(
    conditional_mean_fn: Callable[[Float[Array, " D"]], Float[Array, " M"]],
    conditional_var_fn: Callable[[Float[Array, " D"]], Float[Array, "M M"]],
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> SLRResult:
    r"""Moment-match a nonlinear conditional to a linear-Gaussian surrogate.

    Given a conditional ``p(y | f)`` described by its first two moments and
    a Gaussian state ``q(f) = N(m_q, C_ff)``, returns the linear-Gaussian
    model ``p(y | f) ≈ N(y | A f + b, Ω)`` whose moments match those of the
    true conditional under ``q``.

    With

    $$
    \mu = \mathbb{E}_q[\mathbb{E}[y \mid f]], \quad
    C_{xf} = \mathrm{Cov}_q(f, \mathbb{E}[y \mid f]), \quad
    S = \mathrm{Var}_q(\mathbb{E}[y \mid f])
        + \mathbb{E}_q[\mathrm{Var}[y \mid f]],
    $$

    the surrogate is

    $$
    A = C_{xf}^\top C_{ff}^{-1}, \quad
    b = \mu - A m_q, \quad
    \Omega = S - A C_{xf}.
    $$

    This is the shared primitive behind posterior linearisation, the
    iterated Kalman smoother (IPLF / IPLS), and the conjugate Gaussian site
    that an EP cavity meets. Feeding ``(A, b, omega)`` into a standard
    linear-Gaussian Kalman update reproduces posterior-linearisation
    inference exactly.

    All expectations are evaluated with ``integrator``, so the accuracy /
    cost trade-off is chosen by the caller. Only the conditional moments
    are needed — no autodiff runs through the quadrature.

    Args:
        conditional_mean_fn: Maps a latent ``f`` of shape ``(D,)`` to the
            conditional mean ``E[y | f]`` of shape ``(M,)``.
        conditional_var_fn: Maps a latent ``f`` of shape ``(D,)`` to the
            conditional covariance ``Var[y | f]``. Either the full
            ``(M, M)`` covariance or, for conditionally independent
            outputs, an ``(M,)`` vector of variances.
        state: Gaussian state ``q(f)``. Its covariance must be symmetric.
        integrator: Point-based rule used for the expectations, e.g.
            `gaussx.FifthOrderCubatureIntegrator`.

    Returns:
        `SLRResult` with the gain, offset, residual covariance, and
        expected conditional mean.

    Raises:
        NotImplementedError: If ``integrator`` is not point-based.
        ValueError: If ``conditional_var_fn`` does not return an ``(M,)``
            or ``(M, M)`` array matching the ``M`` outputs of
            ``conditional_mean_fn``.
    """
    chi, w_m, w_c = integrator.points_and_weights(state)

    # First pass: moments of the conditional mean under q(f).
    G = jax.vmap(conditional_mean_fn)(chi)  # (P, M)
    moments = assemble_propagation_result(chi, G, state.mean, w_m, w_c)
    mu = moments.state.mean  # (M,)
    var_of_mean = moments.state.cov.as_matrix()  # Var_q(E[y|f])
    cross_cov = moments.cross_cov  # C_xf, (D, M)
    assert cross_cov is not None, "Integrator must return cross_cov"

    # Second pass: E_q[Var[y|f]], reusing the same points. The trailing
    # shape is checked against M rather than just the rank: a (1,) or
    # (1, 1) return would otherwise broadcast across the whole of ``S``
    # and silently corrupt every entry of ``omega``.
    out_dim = G.shape[-1]
    V = jax.vmap(conditional_var_fn)(chi)
    if V.shape[1:] == (out_dim,):  # (P, M) diagonal variances
        mean_of_var = jnp.diag(jnp.sum(w_m[:, None] * V, axis=0))
    elif V.shape[1:] == (out_dim, out_dim):  # (P, M, M) full covariances
        mean_of_var = jnp.sum(w_m[:, None, None] * V, axis=0)
    else:
        msg = (
            f"conditional_var_fn must return an ({out_dim},) or "
            f"({out_dim}, {out_dim}) array per point to match the "
            f"{out_dim} outputs of conditional_mean_fn, got trailing "
            f"shape {V.shape[1:]}."
        )
        raise ValueError(msg)

    S = var_of_mean + mean_of_var

    # A = C_xf^T C_ff^{-1}; solve_rows solves C_ff x = row, and C_ff is
    # symmetric, so stacking the rows of C_xf^T gives A directly.
    A = solve_rows(state.cov, cross_cov.T)  # (M, D)

    b = mu - A @ state.mean
    omega = symmetrize(S - A @ cross_cov)

    return SLRResult(A=A, b=b, omega=omega, mu=mu)

Hyperparameter designs (INLA)

theta_design turns a hyperparameter log-posterior \(\log\tilde\pi(\theta \mid y)\) and its mode \(\theta^\ast\) into a small set of integration points \(\theta_k\) and normalised log-weights, the way R-INLA integrates over \(\theta\) (Rue, Martino & Chopin, 2009, §6.5). Points are laid out in the standardised coordinates \(\theta(z) = \theta^\ast + V\Lambda^{-1/2}z\), where \(-\nabla^2\log\tilde\pi(\theta^\ast) = V\Lambda V^\top\) (exact, from jax.hessian), and each design weight is corrected by the evaluated log-posterior:

• "eb" — empirical Bayes, the mode alone;
• "grid" — an axis-wise grid in \(z\), truncated grid_threshold log-units below the mode (the default for \(m \le 2\));
• "ccd" — a central composite design: centre, \(2m\) axial points and a resolution-V fractional factorial on the sphere of radius \(f_0\sqrt m\), \(O(m^2)\) points (15 at \(m = 3\), 27 at \(m = 5\)) against \(3^m\) for a grid (the default for \(m > 2\)).

The design is a pure function of any callable, so mode-finding stays with the caller. A typical INLA step fits one Laplace approximation per design point in a single vmap (laplace_mode lands with G8):

pts, logw = gaussx.theta_design(log_post, theta_star, method="ccd")  # (15, 3), (15,)
fits = jax.vmap(lambda th: gaussx.laplace_mode(prior_at(th), lik, y))(pts)
post_mean = einx.dot("k, k n -> n", jnp.exp(logw), fits.mode)

Structured linear algebra and Gaussian primitives for JAX.

theta_design(log_post: Callable[[Array], Array], mode: Float[Array, ' m'], *, method: Literal['eb', 'grid', 'ccd'] | None = None, hessian: Float[Array, 'm m'] | None = None, grid_step: float = 1.0, grid_threshold: float = 2.5, ccd_f0: float = 1.1) -> tuple[Float[Array, 'K m'], Float[Array, ' K']]

Integration points and log-weights over hyperparameters \(\theta\).

Builds the design R-INLA uses to integrate over \(\theta\). With \(-\nabla^2 \log\tilde\pi(\theta^\ast \mid y) = V \Lambda V^\top\) at the mode \(\theta^\ast\), points are placed in the standardised coordinates \(z\) of

\[ \theta(z) = \theta^\ast + V \Lambda^{-1/2} z , \]

in which the posterior is approximately \(\mathcal N(0, I)\). Each point gets a design weight \(\Delta_k\), which is corrected by the evaluated posterior, \(\log w_k = \log\Delta_k + \log\tilde\pi(\theta_k \mid y)\), and normalised so that \(\sum_k w_k = 1\). Then \(\tilde\pi(x_i \mid y) \approx \sum_k w_k\,\tilde\pi(x_i \mid \theta_k, y)\).

Methods:

• "eb": empirical Bayes, the mode alone.
• "grid": walk each \(z\)-axis in both directions in steps of grid_step while \(\log\tilde\pi\) stays within grid_threshold log-units of the mode, take the product grid of those axis ranges, and keep its points that are within grid_threshold of the mode. \(\Delta_k\) is constant. The point count depends on log_post, so this method cannot be traced under jax.jit.
• "ccd": a central composite design, \(O(m^2)\) points: the centre, \(2m\) axial points and a resolution-V fractional factorial \(2^{m-p}\), all non-centre points on the sphere of radius \(f_0\sqrt m\) (15 points at \(m = 3\), 27 at \(m = 5\)). The centre and shell weights are fixed, as in Rue et al. (2009) §6.5, so that the design integrates a Gaussian posterior exactly up to its second moments: \(\Delta_0 \propto 1 - f_0^{-2}\) and \(\Delta_{\text{shell}} \propto e^{m f_0^2/2} / ((K - 1) f_0^2)\).

Parameters:

Name	Type	Description	Default
log_post	Callable[[Array], Array]	Unnormalised log-posterior \(\log\tilde\pi(\theta \mid y)\), mapping a (m,) array to a scalar.	required
mode	Float[Array, ' m']	Its mode \(\theta^\ast\), shape (m,). Finding it is the caller's job.	required
method	Literal['eb', 'grid', 'ccd'] | None	"eb", "grid" or "ccd". None (default) picks "grid" for m <= 2 and "ccd" for m > 2.	None
hessian	Float[Array, 'm m'] | None	Hessian of log_post at mode, shape (m, m) (negative definite). If None, computed reverse-over-reverse (jax.jacrev(jax.jacrev(log_post))), which also differentiates through custom-VJP-only code such as the sparse Cholesky log-determinant (jax.hessian is forward-over-reverse).	None
grid_step	float	Step size in \(z\) for "grid".	1.0
grid_threshold	float	Log-density drop from the mode beyond which "grid" stops and discards points.	2.5
ccd_f0	float	Shell radius factor \(f_0 > 1\) for "ccd".	1.1

Returns:

Type	Description
Float[Array, 'K m']	Tuple (points, log_weights): points \(\theta_k\) of shape
Float[Array, ' K']	(K, m) and normalised log-weights of shape (K,)
tuple[Float[Array, 'K m'], Float[Array, ' K']]	(logsumexp(log_weights) == 0). The first point is the mode.

Example

def log_post(theta):
    return -0.5 * jnp.sum((theta - 1.0) ** 2 / jnp.array([1.0, 4.0, 9.0]))


pts, logw = gaussx.theta_design(log_post, jnp.ones(3), method="ccd")
pts.shape, logw.shape  # (15, 3), (15,)
post_mean = einx.dot("k, k m -> m", jnp.exp(logw), pts)  # == 1.0

# In INLA, one batched Laplace fit per design point:
fits = jax.vmap(lambda th: gaussx.laplace_mode(prior_at(th), lik, y))(pts)
x_mean = einx.dot("k, k n -> n", jnp.exp(logw), fits.mode)

Source code in src/gaussx/_quadrature/_theta_design.py

def theta_design(
    log_post: Callable[[Array], Array],
    mode: Float[Array, " m"],
    *,
    method: Literal["eb", "grid", "ccd"] | None = None,
    hessian: Float[Array, "m m"] | None = None,
    grid_step: float = 1.0,
    grid_threshold: float = 2.5,
    ccd_f0: float = 1.1,
) -> tuple[Float[Array, "K m"], Float[Array, " K"]]:
    r"""Integration points and log-weights over hyperparameters $\theta$.

    Builds the design R-INLA uses to integrate over $\theta$. With
    $-\nabla^2 \log\tilde\pi(\theta^\ast \mid y) = V \Lambda V^\top$ at the
    mode $\theta^\ast$, points are placed in the standardised coordinates
    $z$ of

    $$
    \theta(z) = \theta^\ast + V \Lambda^{-1/2} z ,
    $$

    in which the posterior is approximately $\mathcal N(0, I)$. Each point
    gets a design weight $\Delta_k$, which is corrected by the evaluated
    posterior, $\log w_k = \log\Delta_k + \log\tilde\pi(\theta_k \mid y)$,
    and normalised so that $\sum_k w_k = 1$. Then
    $\tilde\pi(x_i \mid y) \approx \sum_k w_k\,\tilde\pi(x_i \mid \theta_k, y)$.

    Methods:

    - `"eb"`: empirical Bayes, the mode alone.
    - `"grid"`: walk each $z$-axis in both directions in steps of
      `grid_step` while $\log\tilde\pi$ stays within `grid_threshold`
      log-units of the mode, take the product grid of those axis ranges,
      and keep its points that are within `grid_threshold` of the mode.
      $\Delta_k$ is constant. The point count depends on `log_post`, so
      this method cannot be traced under `jax.jit`.
    - `"ccd"`: a central composite design, $O(m^2)$ points: the centre,
      $2m$ axial points and a resolution-V fractional factorial
      $2^{m-p}$, all non-centre points on the sphere of radius
      $f_0\sqrt m$ (15 points at $m = 3$, 27 at $m = 5$). The centre and
      shell weights are fixed, as in Rue et al. (2009) §6.5, so that the
      design integrates a Gaussian posterior exactly up to its second
      moments:
      $\Delta_0 \propto 1 - f_0^{-2}$ and
      $\Delta_{\text{shell}} \propto e^{m f_0^2/2} / ((K - 1) f_0^2)$.

    Args:
        log_post: Unnormalised log-posterior $\log\tilde\pi(\theta \mid y)$,
            mapping a ``(m,)`` array to a scalar.
        mode: Its mode $\theta^\ast$, shape ``(m,)``. Finding it is the
            caller's job.
        method: ``"eb"``, ``"grid"`` or ``"ccd"``. ``None`` (default)
            picks ``"grid"`` for ``m <= 2`` and ``"ccd"`` for ``m > 2``.
        hessian: Hessian of `log_post` at `mode`, shape ``(m, m)``
            (negative definite). If ``None``, computed reverse-over-reverse
            (``jax.jacrev(jax.jacrev(log_post))``), which also differentiates
            through custom-VJP-only code such as the sparse Cholesky
            log-determinant (`jax.hessian` is forward-over-reverse).
        grid_step: Step size in $z$ for ``"grid"``.
        grid_threshold: Log-density drop from the mode beyond which
            ``"grid"`` stops and discards points.
        ccd_f0: Shell radius factor $f_0 > 1$ for ``"ccd"``.

    Returns:
        Tuple ``(points, log_weights)``: points $\theta_k$ of shape
        ``(K, m)`` and normalised log-weights of shape ``(K,)``
        (``logsumexp(log_weights) == 0``). The first point is the mode.

    Example:
        ```python
        def log_post(theta):
            return -0.5 * jnp.sum((theta - 1.0) ** 2 / jnp.array([1.0, 4.0, 9.0]))


        pts, logw = gaussx.theta_design(log_post, jnp.ones(3), method="ccd")
        pts.shape, logw.shape  # (15, 3), (15,)
        post_mean = einx.dot("k, k m -> m", jnp.exp(logw), pts)  # == 1.0

        # In INLA, one batched Laplace fit per design point:
        fits = jax.vmap(lambda th: gaussx.laplace_mode(prior_at(th), lik, y))(pts)
        x_mean = einx.dot("k, k n -> n", jnp.exp(logw), fits.mode)
        ```
    """
    mode = jnp.asarray(mode)
    if mode.ndim != 1:
        raise ValueError(f"mode must be 1-D, got shape {mode.shape}")
    m = mode.shape[0]
    if method is None:
        method = "grid" if m <= 2 else "ccd"
    dtype = mode.dtype

    if method == "eb":
        return rearrange(mode, "m -> 1 m"), jnp.zeros((1,), dtype=dtype)

    if hessian is None:
        hessian = jax.jacrev(jax.jacrev(log_post))(mode)
    evals, evecs = jnp.linalg.eigh(-jnp.asarray(hessian, dtype=dtype))
    # theta(z) = mode + V Lambda^{-1/2} z
    scale = einx.multiply("i j, j -> i j", evecs, jax.lax.rsqrt(evals))

    def to_theta(z: Float[Array, "K m"]) -> Float[Array, "K m"]:
        return einx.add("k i, i -> k i", einsum(scale, z, "i j, k j -> k i"), mode)

    if method == "grid":
        if grid_step <= 0 or grid_threshold <= 0:
            raise ValueError("grid_step and grid_threshold must be positive")
        z = _grid_z(log_post, scale, mode, grid_step, grid_threshold)
        points = to_theta(jnp.asarray(z, dtype=dtype))
        lp = jax.vmap(log_post)(points)
        keep = np.asarray(lp >= log_post(mode) - grid_threshold)
        points, lp = points[keep], lp[keep]
        return points, lp - jsp.logsumexp(lp)

    if method == "ccd":
        if ccd_f0 <= 1:
            raise ValueError(f"ccd_f0 must be > 1, got {ccd_f0}")
        z_unit = _ccd_design(m)  # centre first; shell points on radius sqrt(m)
        n_points = z_unit.shape[0]
        log_delta = np.full(
            n_points,
            m * ccd_f0**2 / 2 - math.log((n_points - 1) * ccd_f0**2),
        )
        log_delta[0] = math.log1p(-(ccd_f0**-2))
        points = to_theta(jnp.asarray(ccd_f0 * z_unit, dtype=dtype))
        logw = jnp.asarray(log_delta, dtype=dtype) + jax.vmap(log_post)(points)
        return points, logw - jsp.logsumexp(logw)

    raise ValueError(f"method must be 'eb', 'grid' or 'ccd', got {method!r}")

Kernel expectations & uncertain-input GPs

The \(\Psi\)-statistics \(\Psi_0 = \mathbb{E}[k(x,x)]\), \(\Psi_1 = \mathbb{E}[k(x, X)]\), \(\Psi_2 = \mathbb{E}[k(x,\cdot)k(x,\cdot)^\top]\) and the GP / SVGP / VGP / BGPLVM predictive equations for inputs that are themselves Gaussian.

Structured linear algebra and Gaussian primitives for JAX.

AnalyticalPsiStatistics

Bases: Protocol

Protocol for kernels with closed-form Ψ statistics.

Ψ statistics are required for uncertain-input GP models (e.g., BGPLVM). A kernel implementing this protocol provides analytical formulae instead of requiring numerical integration.

Source code in src/gaussx/_quadrature/_psi_statistics.py

@runtime_checkable
class AnalyticalPsiStatistics(Protocol):
    """Protocol for kernels with closed-form Ψ statistics.

    Ψ statistics are required for uncertain-input GP models
    (e.g., BGPLVM). A kernel implementing this protocol provides
    analytical formulae instead of requiring numerical integration.
    """

    def psi0(self, state: GaussianState) -> Float[Array, ""]:
        """Compute Ψ₀ = E[k(x, x)] (scalar)."""
        ...

    def psi1(
        self,
        state: GaussianState,
        X_train: Float[Array, "M D"],
    ) -> Float[Array, " M"]:
        """Compute Ψ₁ᵢ = E[k(x, xᵢ)], shape ``(M,)``."""
        ...

    def psi2(
        self,
        state: GaussianState,
        X_train: Float[Array, "M D"],
    ) -> Float[Array, "M M"]:
        """Compute Ψ₂ᵢⱼ = E[k(x, xᵢ) k(x, xⱼ)], shape ``(M, M)``."""
        ...

psi0(state: GaussianState) -> Float[Array, '']

Compute Ψ₀ = E[k(x, x)] (scalar).

Source code in src/gaussx/_quadrature/_psi_statistics.py

def psi0(self, state: GaussianState) -> Float[Array, ""]:
    """Compute Ψ₀ = E[k(x, x)] (scalar)."""
    ...

psi1(state: GaussianState, X_train: Float[Array, 'M D']) -> Float[Array, ' M']

Compute Ψ₁ᵢ = E[k(x, xᵢ)], shape (M,).

Source code in src/gaussx/_quadrature/_psi_statistics.py

def psi1(
    self,
    state: GaussianState,
    X_train: Float[Array, "M D"],
) -> Float[Array, " M"]:
    """Compute Ψ₁ᵢ = E[k(x, xᵢ)], shape ``(M,)``."""
    ...

psi2(state: GaussianState, X_train: Float[Array, 'M D']) -> Float[Array, 'M M']

Compute Ψ₂ᵢⱼ = E[k(x, xᵢ) k(x, xⱼ)], shape (M, M).

Source code in src/gaussx/_quadrature/_psi_statistics.py

def psi2(
    self,
    state: GaussianState,
    X_train: Float[Array, "M D"],
) -> Float[Array, "M M"]:
    """Compute Ψ₂ᵢⱼ = E[k(x, xᵢ) k(x, xⱼ)], shape ``(M, M)``."""
    ...

kernel_expectations(kernel_fn: Callable[[Float[Array, ' D'], Float[Array, ' D']], Float[Array, '']], state: GaussianState, X_train: Float[Array, 'N_train D'], integrator: AbstractIntegrator) -> tuple[Float[Array, ''], Float[Array, ' N_train'], Float[Array, 'N_train N_train']]

Compute kernel expectations Psi_0, Psi_1, Psi_2 for uncertain inputs.

These are the core quantities for GP inference with uncertain inputs:

Psi_0 = E[k(x, x)]                      scalar
Psi_1_i = E[k(x, x_i)]                  (N_train,)
Psi_2_{ij} = E[k(x, x_i) k(x, x_j)]    (N_train, N_train)

For a point-based rule with points chi_p and mean weights w_p they are computed directly from the (P, N_train) kernel evaluations K_{pi} = k(chi_p, x_i):

\[ \Psi_0 = \sum_p w_p\, k(\chi_p, \chi_p), \quad \Psi_{1,i} = \sum_p w_p K_{pi}, \quad \Psi_{2,ij} = \sum_p w_p K_{pi} K_{pj}, \]

so no (N_train^2, N_train^2) output covariance or per-point outer product is ever built. Integrators without points (e.g. Taylor) go through mean_expectation on the flattened outer product.

Parameters:

Name	Type	Description	Default
kernel_fn	Callable[[Float[Array, ' D'], Float[Array, ' D']], Float[Array, '']]	Kernel function k(x, x') -> scalar.	required
state	GaussianState	Uncertain input distribution x ~ N(mu, Sigma).	required
X_train	Float[Array, 'N_train D']	Training points, shape (N_train, D).	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, ' N_train'], Float[Array, 'N_train N_train']]	Tuple (Psi_0, Psi_1, Psi_2).

Source code in src/gaussx/_quadrature/_gp_predict.py

def kernel_expectations(
    kernel_fn: Callable[[Float[Array, " D"], Float[Array, " D"]], Float[Array, ""]],
    state: GaussianState,
    X_train: Float[Array, "N_train D"],
    integrator: AbstractIntegrator,
) -> tuple[Float[Array, ""], Float[Array, " N_train"], Float[Array, "N_train N_train"]]:
    r"""Compute kernel expectations Psi_0, Psi_1, Psi_2 for uncertain inputs.

    These are the core quantities for GP inference with uncertain inputs:

        Psi_0 = E[k(x, x)]                      scalar
        Psi_1_i = E[k(x, x_i)]                  (N_train,)
        Psi_2_{ij} = E[k(x, x_i) k(x, x_j)]    (N_train, N_train)

    For a point-based rule with points ``chi_p`` and mean weights ``w_p``
    they are computed directly from the ``(P, N_train)`` kernel evaluations
    ``K_{pi} = k(chi_p, x_i)``:

    $$
    \Psi_0 = \sum_p w_p\, k(\chi_p, \chi_p), \quad
    \Psi_{1,i} = \sum_p w_p K_{pi}, \quad
    \Psi_{2,ij} = \sum_p w_p K_{pi} K_{pj},
    $$

    so no ``(N_train^2, N_train^2)`` output covariance or per-point outer
    product is ever built. Integrators without points (e.g. Taylor) go
    through `mean_expectation` on the flattened outer product.

    Args:
        kernel_fn: Kernel function ``k(x, x') -> scalar``.
        state: Uncertain input distribution ``x ~ N(mu, Sigma)``.
        X_train: Training points, shape ``(N_train, D)``.
        integrator: Integration method.

    Returns:
        Tuple ``(Psi_0, Psi_1, Psi_2)``.
    """

    def k_row(x: Float[Array, " D"]) -> Float[Array, " N_train"]:
        return jax.vmap(lambda xi: kernel_fn(x, xi))(X_train)

    rule = _points_and_weights_or_none(integrator, state)
    if rule is not None:
        chi, w, _ = rule
        K_self = jax.vmap(lambda x: kernel_fn(x, x))(chi)  # (P,)
        Kx = jax.vmap(k_row)(chi)  # (P, N_train)
        Psi_0 = einsum(w, K_self, "p, p ->")
        Psi_1 = einsum(w, Kx, "p, p i -> i")
        Psi_2 = einsum(einx.multiply("p, p i -> p i", w, Kx), Kx, "p i, p j -> i j")
        return Psi_0, Psi_1, Psi_2

    # Psi_0 = E[k(x, x)]
    Psi_0 = mean_expectation(
        lambda x: jnp.atleast_1d(kernel_fn(x, x)),
        state,
        integrator,
    )[0]

    # Psi_1_i = E[k(x, x_i)]
    Psi_1 = mean_expectation(k_row, state, integrator)

    # Psi_2_{ij} = E[k(x, x_i) * k(x, x_j)]
    N_train = X_train.shape[0]

    def psi2_flat_fn(x: Float[Array, " D"]) -> Float[Array, " NN"]:
        k_vec = k_row(x)
        return rearrange(jnp.outer(k_vec, k_vec), "i j -> (i j)")

    Psi_2_flat = mean_expectation(psi2_flat_fn, state, integrator)
    Psi_2 = rearrange(Psi_2_flat, "(i j) -> i j", i=N_train, j=N_train)

    return Psi_0, Psi_1, Psi_2

compute_psi_statistics(kernel: object, state: GaussianState, X_train: Float[Array, 'M D'], *, integrator: AbstractIntegrator | None = None) -> tuple[Float[Array, ''], Float[Array, ' M'], Float[Array, 'M M']]

Compute Ψ statistics, dispatching to analytical or numerical.

If kernel implements AnalyticalPsiStatistics, uses the closed-form methods. Otherwise, falls back to numerical integration via the provided integrator:

Ψ₀   = E[k(x, x)]                   scalar
Ψ₁ᵢ  = E[k(x, xᵢ)]                 (M,)
Ψ₂ᵢⱼ = E[k(x, xᵢ) k(x, xⱼ)]       (M, M)

Parameters:

Name	Type	Description	Default
kernel	object	Kernel object, optionally implementing AnalyticalPsiStatistics.	required
state	GaussianState	Input Gaussian distribution x ~ 𝒩(μ, Σ).	required
X_train	Float[Array, 'M D']	Training/inducing points, shape (M, D).	required
integrator	AbstractIntegrator | None	Numerical integrator for fallback. Required if kernel does not implement analytical Ψ statistics.	None

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, ' M'], Float[Array, 'M M']]	Tuple (Ψ₀, Ψ₁, Ψ₂) of Psi statistics.

Raises:

Type	Description
ValueError	If kernel has no analytical Ψ statistics and no integrator is provided.

Source code in src/gaussx/_quadrature/_psi_statistics.py

def compute_psi_statistics(
    kernel: object,
    state: GaussianState,
    X_train: Float[Array, "M D"],
    *,
    integrator: AbstractIntegrator | None = None,
) -> tuple[Float[Array, ""], Float[Array, " M"], Float[Array, "M M"]]:
    """Compute Ψ statistics, dispatching to analytical or numerical.

    If ``kernel`` implements `AnalyticalPsiStatistics`, uses
    the closed-form methods. Otherwise, falls back to numerical
    integration via the provided integrator:

        Ψ₀   = E[k(x, x)]                   scalar
        Ψ₁ᵢ  = E[k(x, xᵢ)]                 (M,)
        Ψ₂ᵢⱼ = E[k(x, xᵢ) k(x, xⱼ)]       (M, M)

    Args:
        kernel: Kernel object, optionally implementing
            `AnalyticalPsiStatistics`.
        state: Input Gaussian distribution x ~ 𝒩(μ, Σ).
        X_train: Training/inducing points, shape ``(M, D)``.
        integrator: Numerical integrator for fallback. Required if
            ``kernel`` does not implement analytical Ψ statistics.

    Returns:
        Tuple ``(Ψ₀, Ψ₁, Ψ₂)`` of Psi statistics.

    Raises:
        ValueError: If ``kernel`` has no analytical Ψ statistics
            and no integrator is provided.
    """
    if isinstance(kernel, AnalyticalPsiStatistics):
        psi0 = kernel.psi0(state)
        psi1 = kernel.psi1(state, X_train)
        psi2 = kernel.psi2(state, X_train)
        return psi0, psi1, psi2

    if integrator is None:
        msg = (
            "Kernel does not implement AnalyticalPsiStatistics and no "
            "integrator was provided. Either implement the protocol on "
            "the kernel or pass an integrator for numerical computation."
        )
        raise ValueError(msg)

    # ── Numerical fallback: one implementation, `kernel_expectations` ──

    return kernel_expectations(cast(Callable, kernel), state, X_train, integrator)

uncertain_gp_predict(kernel_fn: Callable, X_train: Float[Array, 'N_train D'], alpha: Float[Array, ' N_train'], K_inv: lx.AbstractLinearOperator, state: GaussianState, integrator: AbstractIntegrator) -> tuple[Float[Array, ''], Float[Array, '']]

Predictive mean and variance for GP with uncertain inputs.

Uses kernel expectations:

mu_pred = Psi_1 @ alpha
var_pred = Psi_0 - tr(K_inv @ Psi_2) + alpha^T @ Psi_2 @ alpha - mu_pred^2

Parameters:

Name	Type	Description	Default
kernel_fn	Callable	Kernel function k(x, x') -> scalar.	required
X_train	Float[Array, 'N_train D']	Training points, shape (N_train, D).	required
alpha	Float[Array, ' N_train']	Precomputed weights K^{-1} y, shape (N_train,).	required
K_inv	AbstractLinearOperator	Inverse of training kernel matrix operator.	required
state	GaussianState	Uncertain test input x ~ N(mu, Sigma).	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, '']]	Tuple (mean, variance) — scalar predictive moments.

Source code in src/gaussx/_quadrature/_gp_predict.py

def uncertain_gp_predict(
    kernel_fn: Callable,
    X_train: Float[Array, "N_train D"],
    alpha: Float[Array, " N_train"],
    K_inv: lx.AbstractLinearOperator,
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    r"""Predictive mean and variance for GP with uncertain inputs.

    Uses kernel expectations:

        mu_pred = Psi_1 @ alpha
        var_pred = Psi_0 - tr(K_inv @ Psi_2) + alpha^T @ Psi_2 @ alpha - mu_pred^2

    Args:
        kernel_fn: Kernel function ``k(x, x') -> scalar``.
        X_train: Training points, shape ``(N_train, D)``.
        alpha: Precomputed weights ``K^{-1} y``, shape ``(N_train,)``.
        K_inv: Inverse of training kernel matrix operator.
        state: Uncertain test input ``x ~ N(mu, Sigma)``.
        integrator: Integration method.

    Returns:
        Tuple ``(mean, variance)`` — scalar predictive moments.
    """

    Psi_0, Psi_1, Psi_2 = kernel_expectations(kernel_fn, state, X_train, integrator)

    mu_pred = jnp.dot(Psi_1, alpha)

    Psi_2_op = lx.MatrixLinearOperator(Psi_2)
    tr_term = trace_product(K_inv, Psi_2_op)
    quad_term = alpha @ Psi_2 @ alpha

    var_pred = Psi_0 - tr_term + quad_term - mu_pred**2
    var_pred = jnp.clip(var_pred, 0.0)

    return mu_pred, var_pred

uncertain_gp_predict_mc(predict_fn: Callable[[Float[Array, ' D']], tuple[Float[Array, ''], Float[Array, '']]], state: GaussianState, n_particles: int = 100, key: jax.Array | None = None) -> tuple[Float[Array, ''], Float[Array, '']]

Monte Carlo GP prediction with uncertain inputs.

Alternative to analytic kernel expectations when Psi integrals are intractable. Uses law of total variance:

mu = mean(particle_means)
var = var(particle_means) + mean(particle_vars)

Parameters:

Name	Type	Description	Default
predict_fn	Callable[[Float[Array, ' D']], tuple[Float[Array, ''], Float[Array, '']]]	GP predictor mapping (D,) -> (mean, var).	required
state	GaussianState	Uncertain test input x ~ N(mu, Sigma).	required
n_particles	int	Number of Monte Carlo particles. Default 100.	100
key	Array | None	PRNG key. If None, uses jax.random.key(0).	None

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, '']]	Tuple (mean, variance) — scalar predictive moments.

Source code in src/gaussx/_quadrature/_gp_predict.py

def uncertain_gp_predict_mc(
    predict_fn: Callable[
        [Float[Array, " D"]], tuple[Float[Array, ""], Float[Array, ""]]
    ],
    state: GaussianState,
    n_particles: int = 100,
    key: jax.Array | None = None,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    r"""Monte Carlo GP prediction with uncertain inputs.

    Alternative to analytic kernel expectations when Psi integrals are
    intractable. Uses law of total variance:

        mu = mean(particle_means)
        var = var(particle_means) + mean(particle_vars)

    Args:
        predict_fn: GP predictor mapping ``(D,) -> (mean, var)``.
        state: Uncertain test input ``x ~ N(mu, Sigma)``.
        n_particles: Number of Monte Carlo particles. Default ``100``.
        key: PRNG key. If ``None``, uses ``jax.random.key(0)``.

    Returns:
        Tuple ``(mean, variance)`` — scalar predictive moments.
    """
    mu = state.mean
    N = mu.shape[0]

    if key is None:
        key = jr.key(0)

    # Sample inputs
    L = cholesky(state.cov).as_matrix()
    eps = jr.normal(key, (n_particles, N), dtype=jnp.result_type(mu, L))
    x_samples = mu[None, :] + eps @ L.T

    # Predict at each sample
    means, variances = jax.vmap(predict_fn)(x_samples)

    # Law of total variance
    pred_mean = jnp.mean(means)
    pred_var = jnp.var(means) + jnp.mean(variances)

    return pred_mean, pred_var

uncertain_svgp_predict(kernel_fn: Callable, Z: Float[Array, 'M D'], alpha: Float[Array, ' M'], Q: lx.AbstractLinearOperator, state: GaussianState, integrator: AbstractIntegrator) -> tuple[Float[Array, ''], Float[Array, '']]

Predictive mean and variance for SVGP with uncertain inputs.

Uses kernel expectations with inducing points:

mu_pred = Psi_1 @ alpha
var_pred = Psi_0 + tr(Q @ Psi_2) + alpha^T @ Psi_2 @ alpha - mu_pred^2

where Q = K_{zz}^{-1} S K_{zz}^{-1} - K_{zz}^{-1} is the variance adjustment operator (see gaussx.svgp_variance_adjustment). The exact uncertain GP trace correction is recovered by setting Q = -K_{zz}^{-1}; if Z equals the training inputs, this matches gaussx.uncertain_gp_predict.

Parameters:

Name	Type	Description	Default
kernel_fn	Callable	Kernel function k(x, x') -> scalar.	required
Z	Float[Array, 'M D']	Inducing points, shape (M, D).	required
alpha	Float[Array, ' M']	Effective weights, shape (M,).	required
Q	AbstractLinearOperator	Variance adjustment operator K_{zz}^{-1} S K_{zz}^{-1} - K_{zz}^{-1}, shape (M, M).	required
state	GaussianState	Uncertain test input x ~ N(mu, Sigma).	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, '']]	Tuple (mean, variance) — scalar predictive moments.

Source code in src/gaussx/_quadrature/_gp_predict.py

def uncertain_svgp_predict(
    kernel_fn: Callable,
    Z: Float[Array, "M D"],
    alpha: Float[Array, " M"],
    Q: lx.AbstractLinearOperator,
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    r"""Predictive mean and variance for SVGP with uncertain inputs.

    Uses kernel expectations with inducing points:

        mu_pred = Psi_1 @ alpha
        var_pred = Psi_0 + tr(Q @ Psi_2) + alpha^T @ Psi_2 @ alpha - mu_pred^2

    where ``Q = K_{zz}^{-1} S K_{zz}^{-1} - K_{zz}^{-1}`` is the variance
    adjustment operator (see `gaussx.svgp_variance_adjustment`).
    The exact uncertain GP trace correction is recovered by setting
    ``Q = -K_{zz}^{-1}``; if ``Z`` equals the training inputs, this matches
    `gaussx.uncertain_gp_predict`.

    Args:
        kernel_fn: Kernel function ``k(x, x') -> scalar``.
        Z: Inducing points, shape ``(M, D)``.
        alpha: Effective weights, shape ``(M,)``.
        Q: Variance adjustment operator ``K_{zz}^{-1} S K_{zz}^{-1} - K_{zz}^{-1}``,
            shape ``(M, M)``.
        state: Uncertain test input ``x ~ N(mu, Sigma)``.
        integrator: Integration method.

    Returns:
        Tuple ``(mean, variance)`` — scalar predictive moments.
    """

    Psi_0, Psi_1, Psi_2 = kernel_expectations(kernel_fn, state, Z, integrator)

    mu_pred = jnp.dot(Psi_1, alpha)

    Psi_2_op = lx.MatrixLinearOperator(Psi_2)
    tr_term = trace_product(Q, Psi_2_op)
    quad_term = alpha @ Psi_2 @ alpha

    var_pred = Psi_0 + tr_term + quad_term - mu_pred**2
    var_pred = jnp.clip(var_pred, 0.0)

    return mu_pred, var_pred

uncertain_vgp_predict(kernel_fn: Callable, X_train: Float[Array, 'N_train D'], alpha: Float[Array, ' N_train'], Q: lx.AbstractLinearOperator, state: GaussianState, integrator: AbstractIntegrator) -> tuple[Float[Array, ''], Float[Array, '']]

Predictive mean and variance for dense VGP with uncertain inputs.

Uses kernel expectations with training points:

mu_pred = Psi_1 @ alpha
var_pred = Psi_0 + tr(Q @ Psi_2) + alpha^T @ Psi_2 @ alpha - mu_pred^2

where Q = K^{-1} S K^{-1} - K^{-1} and alpha = K^{-1} m. The exact uncertain GP is the special case Q = -K^{-1}. By contrast, Q = 0 corresponds to S = K and removes only the trace correction.

Parameters:

Name	Type	Description	Default
kernel_fn	Callable	Kernel function k(x, x') -> scalar.	required
X_train	Float[Array, 'N_train D']	Training points, shape (N_train, D).	required
alpha	Float[Array, ' N_train']	Precomputed weights K^{-1} m, shape (N_train,).	required
Q	AbstractLinearOperator	Variance adjustment operator K^{-1} S K^{-1} - K^{-1}, shape (N_train, N_train).	required
state	GaussianState	Uncertain test input x ~ N(mu, Sigma).	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
tuple[Float[Array, ''], Float[Array, '']]	Tuple (mean, variance) — scalar predictive moments.

Source code in src/gaussx/_quadrature/_gp_predict.py

def uncertain_vgp_predict(
    kernel_fn: Callable,
    X_train: Float[Array, "N_train D"],
    alpha: Float[Array, " N_train"],
    Q: lx.AbstractLinearOperator,
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> tuple[Float[Array, ""], Float[Array, ""]]:
    r"""Predictive mean and variance for dense VGP with uncertain inputs.

    Uses kernel expectations with training points:

        mu_pred = Psi_1 @ alpha
        var_pred = Psi_0 + tr(Q @ Psi_2) + alpha^T @ Psi_2 @ alpha - mu_pred^2

    where ``Q = K^{-1} S K^{-1} - K^{-1}`` and ``alpha = K^{-1} m``.
    The exact uncertain GP is the special case ``Q = -K^{-1}``. By contrast,
    ``Q = 0`` corresponds to ``S = K`` and removes only the trace correction.

    Args:
        kernel_fn: Kernel function ``k(x, x') -> scalar``.
        X_train: Training points, shape ``(N_train, D)``.
        alpha: Precomputed weights ``K^{-1} m``, shape ``(N_train,)``.
        Q: Variance adjustment operator ``K^{-1} S K^{-1} - K^{-1}``,
            shape ``(N_train, N_train)``.
        state: Uncertain test input ``x ~ N(mu, Sigma)``.
        integrator: Integration method.

    Returns:
        Tuple ``(mean, variance)`` — scalar predictive moments.
    """

    Psi_0, Psi_1, Psi_2 = kernel_expectations(kernel_fn, state, X_train, integrator)

    mu_pred = jnp.dot(Psi_1, alpha)

    Psi_2_op = lx.MatrixLinearOperator(Psi_2)
    tr_term = trace_product(Q, Psi_2_op)
    quad_term = alpha @ Psi_2 @ alpha

    var_pred = Psi_0 + tr_term + quad_term - mu_pred**2
    var_pred = jnp.clip(var_pred, 0.0)

    return mu_pred, var_pred

uncertain_bgplvm_predict(kernel_fn: Callable, X_train: Float[Array, 'N_train D'], alpha: Float[Array, 'N_train D_out'], K_inv: lx.AbstractLinearOperator, state: GaussianState, integrator: AbstractIntegrator) -> tuple[Float[Array, ' D_out'], Float[Array, ' D_out']]

Multi-output uncertain GP prediction for BGPLVM.

Maps a latent variable z ~ N(mu, Sigma) to high-dimensional reconstruction y using per-output GP weights:

mu_pred_d = Psi_1 @ alpha_d
var_pred_d = Psi_0 - tr(K_inv @ Psi_2)
           + alpha_d^T @ Psi_2 @ alpha_d - mu_pred_d^2

This intentionally uses the exact GP trace term for every output dimension. Unlike gaussx.uncertain_vgp_predict and gaussx.uncertain_svgp_predict, there is no separate variational covariance correction operator.

Parameters:

Name	Type	Description	Default
kernel_fn	Callable	Kernel function k(x, x') -> scalar.	required
X_train	Float[Array, 'N_train D']	Training points, shape (N_train, D).	required
alpha	Float[Array, 'N_train D_out']	Multi-output weights K^{-1} Y, shape (N_train, D_out).	required
K_inv	AbstractLinearOperator	Inverse of training kernel matrix operator.	required
state	GaussianState	Uncertain latent input z ~ N(mu, Sigma).	required
integrator	AbstractIntegrator	Integration method.	required

Returns:

Type	Description
tuple[Float[Array, ' D_out'], Float[Array, ' D_out']]	Tuple (mean, variance) — predictive moments each of shape (D_out,).

Source code in src/gaussx/_quadrature/_gp_predict.py

def uncertain_bgplvm_predict(
    kernel_fn: Callable,
    X_train: Float[Array, "N_train D"],
    alpha: Float[Array, "N_train D_out"],
    K_inv: lx.AbstractLinearOperator,
    state: GaussianState,
    integrator: AbstractIntegrator,
) -> tuple[Float[Array, " D_out"], Float[Array, " D_out"]]:
    r"""Multi-output uncertain GP prediction for BGPLVM.

    Maps a latent variable ``z ~ N(mu, Sigma)`` to high-dimensional
    reconstruction ``y`` using per-output GP weights:

        mu_pred_d = Psi_1 @ alpha_d
        var_pred_d = Psi_0 - tr(K_inv @ Psi_2)
                   + alpha_d^T @ Psi_2 @ alpha_d - mu_pred_d^2

    This intentionally uses the exact GP trace term for every output dimension.
    Unlike `gaussx.uncertain_vgp_predict` and
    `gaussx.uncertain_svgp_predict`, there is no separate variational
    covariance correction operator.

    Args:
        kernel_fn: Kernel function ``k(x, x') -> scalar``.
        X_train: Training points, shape ``(N_train, D)``.
        alpha: Multi-output weights ``K^{-1} Y``, shape ``(N_train, D_out)``.
        K_inv: Inverse of training kernel matrix operator.
        state: Uncertain latent input ``z ~ N(mu, Sigma)``.
        integrator: Integration method.

    Returns:
        Tuple ``(mean, variance)`` — predictive moments each of shape ``(D_out,)``.
    """

    Psi_0, Psi_1, Psi_2 = kernel_expectations(kernel_fn, state, X_train, integrator)

    # mu_pred_d = Psi_1 @ alpha_d  for each output dimension
    mu_pred = Psi_1 @ alpha  # (D_out,)

    Psi_2_op = lx.MatrixLinearOperator(Psi_2)
    tr_term = trace_product(K_inv, Psi_2_op)

    # quad_term_d = alpha_d^T @ Psi_2 @ alpha_d  for each output
    quad_term = jnp.sum(alpha * (Psi_2 @ alpha), axis=0)  # (D_out,)

    var_pred = Psi_0 - tr_term + quad_term - mu_pred**2
    var_pred = jnp.clip(var_pred, 0.0)

    return mu_pred, var_pred