Distributions & Exponential Family

Layer 2: Gaussian distributions over structured covariance operators, the sugar operations that probabilistic code actually calls, and the exponential-family (natural-parameter) view used by variational and EP-style inference.

Multivariate normal distributions

NumPyro-compatible distributions whose covariance (or precision) is a lineax operator, so sample / log_prob inherit every structured fast path. MultivariateNormalPrecision carries \(\Lambda = \Sigma^{-1}\) directly — the natural home for natural-parameter guides, where materializing \(\Sigma\) would be wasted work. Both require numpyro to be installed.

Both subclass AbstractMultivariateNormal, so code written against the base works for either parameterisation. It provides covariance_operator and precision_operator (one native, the other a lazy inv), numpyro's dense covariance_matrix / precision_matrix / scale_tril, and p.kl(q) for any mix of the two classes:

\[ \mathrm{KL}(p \,\|\, q) = \tfrac12\Big(\operatorname{tr}(\Lambda_q \Sigma_p) + (\mu_q - \mu_p)^\top \Lambda_q (\mu_q - \mu_p) - N + \log\lvert\Sigma_q\rvert - \log\lvert\Sigma_p\rvert\Big), \qquad \Lambda_q = \Sigma_q^{-1}. \]

When q is precision-parameterised, \(\Lambda_q\) is its native operator, so the trace and quadratic terms need no solve. MultivariateNormalPrecision.variance takes diag_inv's exact structured paths (for example, the \(O(N d^3)\) selected inverse of a BlockTriDiag precision) and densifies only an unstructured precision.

With a stochastic solver (e.g. CGSolver()), log_prob and entropy estimate the log-determinant from the strategy's fixed seed, so every call reuses the same probes. That is fine for SVI gradients, but averaging calls reduces no variance and MCMC targets a fixed pseudo-likelihood. To vary the probes, pass solver=KeyedSolver(CGSolver(), key) with a fresh key per step (see Solvers). The functional gaussian_log_prob, gaussian_entropy and kl_standard_normal take key= directly.

Structured linear algebra and Gaussian primitives for JAX.

AbstractMultivariateNormal

Bases: Distribution

Common surface of MultivariateNormal and MultivariateNormalPrecision.

A subclass stores one native operator: the covariance \(\Sigma\) (MultivariateNormal) or the precision \(\Lambda = \Sigma^{-1}\) (MultivariateNormalPrecision). The other one is the lazy gaussx.inv of it. Code written against this base works for either parameterisation:

• covariance_operator / precision_operator: lineax operators, one native and one a lazy inverse.
• covariance_matrix / precision_matrix / scale_tril: dense arrays with numpyro's names and batch broadcasting. They materialise \(N \times N\) matrices, so they are for small problems and interop.
• mean, log_prob, kl.

Requires the numpyro optional extra (pip install "gaussx[numpyro]").

Source code in src/gaussx/_distributions/_mvn_base.py

class AbstractMultivariateNormal(dist.Distribution):
    r"""Common surface of `MultivariateNormal` and `MultivariateNormalPrecision`.

    A subclass stores one *native* operator: the covariance $\Sigma$
    (`MultivariateNormal`) or the precision $\Lambda = \Sigma^{-1}$
    (`MultivariateNormalPrecision`). The other one is the lazy
    `gaussx.inv` of it. Code written against this base works for either
    parameterisation:

    - ``covariance_operator`` / ``precision_operator``: lineax operators,
      one native and one a lazy inverse.
    - ``covariance_matrix`` / ``precision_matrix`` / ``scale_tril``: dense
      arrays with numpyro's names and batch broadcasting. They materialise
      $N \times N$ matrices, so they are for small problems and interop.
    - ``mean``, ``log_prob``, `kl`.

    Requires the ``numpyro`` optional extra
    (``pip install "gaussx[numpyro]"``).
    """

    arg_constraints = {"loc": dist.constraints.real_vector}  # noqa: RUF012
    support = dist.constraints.real_vector
    reparametrized_params = ["loc"]  # noqa: RUF012

    loc: Float[Array, "*batch N"]
    solver: AbstractSolverStrategy

    @property
    def covariance_operator(self) -> lx.AbstractLinearOperator:
        r"""The covariance $\Sigma$ as a lineax operator."""
        raise NotImplementedError

    @property
    def precision_operator(self) -> lx.AbstractLinearOperator:
        r"""The precision $\Lambda = \Sigma^{-1}$ as a lineax operator."""
        raise NotImplementedError

    def _log_prob_single(self, residual: Float[Array, " N"]) -> Float[Array, ""]:
        raise NotImplementedError

    @validate_sample
    def log_prob(self, value: Float[Array, "*batch N"]) -> Float[Array, "*batch"]:
        residual = value - self.loc
        leading_shape = residual.shape[:-1]
        residual_flat = rearrange(residual, "... D -> (...) D")
        log_prob_flat = jax.vmap(self._log_prob_single)(residual_flat)
        return _reshape_batch(log_prob_flat, leading_shape)

    @lazy_property
    def mean(self) -> Float[Array, "*batch N"]:
        return self.loc

    def _broadcast_matrix(self, matrix: Float[Array, "N N"]) -> Array:
        return jnp.broadcast_to(matrix, self.batch_shape + self.event_shape * 2)

    @lazy_property
    def covariance_matrix(self) -> Float[Array, "*batch N N"]:
        r"""Dense $\Sigma$, broadcast over the batch shape."""
        return self._broadcast_matrix(self.covariance_operator.as_matrix())

    @lazy_property
    def precision_matrix(self) -> Float[Array, "*batch N N"]:
        r"""Dense $\Lambda = \Sigma^{-1}$, broadcast over the batch shape."""
        return self._broadcast_matrix(self.precision_operator.as_matrix())

    @lazy_property
    def scale_tril(self) -> Float[Array, "*batch N N"]:
        r"""Dense lower Cholesky factor $L$ of $\Sigma = L L^\top$ (numpyro's name)."""
        return jnp.linalg.cholesky(self.covariance_matrix)

    def kl(self, other: AbstractMultivariateNormal) -> Float[Array, "*batch"]:
        r"""$\mathrm{KL}(\text{self} \,\|\, \text{other})$ in closed form.

        Any mix of covariance and precision parameterisations works. A
        precision-parameterised ``other`` needs no solve for
        $\operatorname{tr}(\Sigma_q^{-1}\Sigma_p)$ or the quadratic term,
        because the inverse of its lazy covariance is its precision. Batch
        shapes broadcast as in numpyro's ``kl_divergence``. Unlike
        ``numpyro.distributions.kl_divergence``, it uses the closed form
        whatever the solver strategies are.

        Args:
            other: The second Gaussian, with the same event shape.

        Returns:
            The KL divergence, with the broadcast batch shape.
        """
        # lazy import, cycle: _distributions._numpyro_kl -> _distributions._mvn_base
        from gaussx._distributions._numpyro_kl import closed_form_kl

        return closed_form_kl(self, other)

covariance_operator: lx.AbstractLinearOperator property

The covariance \(\Sigma\) as a lineax operator.

precision_operator: lx.AbstractLinearOperator property

The precision \(\Lambda = \Sigma^{-1}\) as a lineax operator.

covariance_matrix() -> Float[Array, '*batch N N']

Dense \(\Sigma\), broadcast over the batch shape.

Source code in src/gaussx/_distributions/_mvn_base.py

@lazy_property
def covariance_matrix(self) -> Float[Array, "*batch N N"]:
    r"""Dense $\Sigma$, broadcast over the batch shape."""
    return self._broadcast_matrix(self.covariance_operator.as_matrix())

precision_matrix() -> Float[Array, '*batch N N']

Dense \(\Lambda = \Sigma^{-1}\), broadcast over the batch shape.

Source code in src/gaussx/_distributions/_mvn_base.py

@lazy_property
def precision_matrix(self) -> Float[Array, "*batch N N"]:
    r"""Dense $\Lambda = \Sigma^{-1}$, broadcast over the batch shape."""
    return self._broadcast_matrix(self.precision_operator.as_matrix())

scale_tril() -> Float[Array, '*batch N N']

Dense lower Cholesky factor \(L\) of \(\Sigma = L L^\top\) (numpyro's name).

Source code in src/gaussx/_distributions/_mvn_base.py

@lazy_property
def scale_tril(self) -> Float[Array, "*batch N N"]:
    r"""Dense lower Cholesky factor $L$ of $\Sigma = L L^\top$ (numpyro's name)."""
    return jnp.linalg.cholesky(self.covariance_matrix)

kl(other: AbstractMultivariateNormal) -> Float[Array, '*batch']

\(\mathrm{KL}(\text{self} \,\|\, \text{other})\) in closed form.

Any mix of covariance and precision parameterisations works. A precision-parameterised other needs no solve for \(\operatorname{tr}(\Sigma_q^{-1}\Sigma_p)\) or the quadratic term, because the inverse of its lazy covariance is its precision. Batch shapes broadcast as in numpyro's kl_divergence. Unlike numpyro.distributions.kl_divergence, it uses the closed form whatever the solver strategies are.

Parameters:

Name	Type	Description	Default
other	AbstractMultivariateNormal	The second Gaussian, with the same event shape.	required

Returns:

Type	Description
Float[Array, '*batch']	The KL divergence, with the broadcast batch shape.

Source code in src/gaussx/_distributions/_mvn_base.py

def kl(self, other: AbstractMultivariateNormal) -> Float[Array, "*batch"]:
    r"""$\mathrm{KL}(\text{self} \,\|\, \text{other})$ in closed form.

    Any mix of covariance and precision parameterisations works. A
    precision-parameterised ``other`` needs no solve for
    $\operatorname{tr}(\Sigma_q^{-1}\Sigma_p)$ or the quadratic term,
    because the inverse of its lazy covariance is its precision. Batch
    shapes broadcast as in numpyro's ``kl_divergence``. Unlike
    ``numpyro.distributions.kl_divergence``, it uses the closed form
    whatever the solver strategies are.

    Args:
        other: The second Gaussian, with the same event shape.

    Returns:
        The KL divergence, with the broadcast batch shape.
    """
    # lazy import, cycle: _distributions._numpyro_kl -> _distributions._mvn_base
    from gaussx._distributions._numpyro_kl import closed_form_kl

    return closed_form_kl(self, other)

MultivariateNormal

Bases: AbstractMultivariateNormal

Multivariate normal parameterized by a lineax linear operator.

Covariance-parameterised member of AbstractMultivariateNormal, which supplies the shared accessors (covariance_operator, precision_operator, covariance_matrix, scale_tril, kl, ...). Unlike numpyro.distributions.MultivariateNormal which requires dense arrays, this distribution accepts any lineax.AbstractLinearOperator as its covariance. This enables efficient log-prob, sampling, and entropy for structured covariances (Kronecker, block-diagonal, low-rank, diagonal, etc.) via gaussx structural dispatch.

Requires the numpyro optional extra (pip install "gaussx[numpyro]").

Parameters:

Name	Type	Description	Default
loc	Float[Array, '*batch N']	Mean vector of shape (N,).	required
cov_operator	AbstractLinearOperator	Covariance as a lineax linear operator of shape (N, N).	required
solver	AbstractSolverStrategy | None	Solver strategy for solve and logdet. Defaults to AutoSolver().	None
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from gaussx._distributions import MultivariateNormal
>>> Sigma = lx.MatrixLinearOperator(
...     jnp.eye(3), lx.positive_semidefinite_tag
... )
>>> d = MultivariateNormal(jnp.zeros(3), Sigma)
>>> round(float(d.log_prob(jnp.ones(3))), 4)
-4.2568

Source code in src/gaussx/_distributions/_mvn.py

class MultivariateNormal(AbstractMultivariateNormal):
    """Multivariate normal parameterized by a lineax linear operator.

    Covariance-parameterised member of `AbstractMultivariateNormal`, which
    supplies the shared accessors (``covariance_operator``,
    ``precision_operator``, ``covariance_matrix``, ``scale_tril``, `kl`, ...).
    Unlike ``numpyro.distributions.MultivariateNormal`` which requires
    dense arrays, this distribution accepts any
    ``lineax.AbstractLinearOperator`` as its covariance. This enables
    efficient log-prob, sampling, and entropy for structured covariances
    (Kronecker, block-diagonal, low-rank, diagonal, etc.) via gaussx
    structural dispatch.

    Requires the ``numpyro`` optional extra
    (``pip install "gaussx[numpyro]"``).

    Args:
        loc: Mean vector of shape ``(N,)``.
        cov_operator: Covariance as a lineax linear operator of shape
            ``(N, N)``.
        solver: Solver strategy for ``solve`` and ``logdet``. Defaults
            to ``AutoSolver()``.
        validate_args: Whether to validate input arguments.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from gaussx._distributions import MultivariateNormal
        >>> Sigma = lx.MatrixLinearOperator(
        ...     jnp.eye(3), lx.positive_semidefinite_tag
        ... )
        >>> d = MultivariateNormal(jnp.zeros(3), Sigma)
        >>> round(float(d.log_prob(jnp.ones(3))), 4)
        -4.2568
    """

    pytree_data_fields = ("loc", "cov_operator", "solver")

    def __init__(
        self,
        loc: Float[Array, "*batch N"],
        cov_operator: lx.AbstractLinearOperator,
        solver: AbstractSolverStrategy | None = None,
        *,
        validate_args: bool | None = None,
    ) -> None:
        if solver is None:
            solver = AutoSolver()
        self.loc = loc
        self.cov_operator = cov_operator
        self.solver = solver
        event_shape = loc.shape[-1:]
        batch_shape = loc.shape[:-1]
        super().__init__(
            batch_shape=batch_shape,
            event_shape=event_shape,
            validate_args=validate_args,
        )

    @property
    def covariance_operator(self) -> lx.AbstractLinearOperator:
        """The covariance operator (native; the same object as ``cov_operator``)."""
        return self.cov_operator

    @property
    def precision_operator(self) -> lx.AbstractLinearOperator:
        """The precision: the lazy `gaussx.inv` of ``cov_operator``."""
        return _inv(self.cov_operator)

    def _log_prob_single(self, residual: Float[Array, " N"]) -> Float[Array, ""]:
        return _gaussian_log_prob_residual(
            residual, self.cov_operator, solver=self.solver
        )

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*batch N"]:
        if key is None:
            raise ValueError(
                "PRNG key must be provided to sample from MultivariateNormal."
            )
        # Delegate to sample_mvn, which dispatches on the covariance's
        # structure (Kronecker, low-rank, scalar multiples, ...) and takes a
        # symmetric square root at the dense fallback, so a semi-definite
        # covariance gives exact finite draws instead of a NaN Cholesky.
        num_samples = math.prod(sample_shape)
        loc = jnp.broadcast_to(self.loc, self.batch_shape + self.event_shape)
        if num_samples == 0:
            # sample_mvn needs at least one draw. Trace one abstractly, so an
            # empty sample still gets its shape checks and its dtype.
            one = jax.eval_shape(
                lambda: sample_mvn(loc, self.cov_operator, key=key, num_samples=1)
            )
            return jnp.zeros(sample_shape + one.shape[1:], dtype=one.dtype)
        draws = sample_mvn(loc, self.cov_operator, key=key, num_samples=num_samples)
        return _unflatten_sample_axis(draws, sample_shape)

    @lazy_property
    def variance(self) -> Float[Array, "*batch N"]:
        return jnp.broadcast_to(
            _diag(self.cov_operator), self.batch_shape + self.event_shape
        )

    def entropy(self) -> Float[Array, "*batch"]:
        # The entropy does not depend on loc, but numpyro's contract is one
        # entry per batch element (Independent.entropy sums over them).
        entropy = gaussian_entropy(self.cov_operator, solver=self.solver)
        return jnp.broadcast_to(entropy, self.batch_shape)

covariance_operator: lx.AbstractLinearOperator property

The covariance operator (native; the same object as cov_operator).

precision_operator: lx.AbstractLinearOperator property

The precision: the lazy gaussx.inv of cov_operator.

MultivariateNormalPrecision

Bases: AbstractMultivariateNormal

Multivariate normal parameterized by a precision (inverse covariance) operator.

This is the natural parameterization for many inference algorithms (e.g. message passing, variational inference in natural coordinates). The precision operator Lambda satisfies Lambda = Sigma^{-1}. Precision-parameterised member of AbstractMultivariateNormal, which supplies the shared accessors (covariance_operator, precision_operator, covariance_matrix, scale_tril, kl, ...).

Requires the numpyro optional extra (pip install "gaussx[numpyro]").

Parameters:

Name	Type	Description	Default
loc	Float[Array, '*batch N']	Mean vector of shape (N,).	required
prec_operator	AbstractLinearOperator	Precision matrix as a lineax linear operator of shape (N, N).	required
solver	AbstractSolverStrategy | None	Solver strategy for solve and logdet. Defaults to AutoSolver().	None
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from gaussx._distributions import MultivariateNormalPrecision
>>> Lambda = lx.MatrixLinearOperator(
...     2.0 * jnp.eye(3), lx.positive_semidefinite_tag
... )
>>> d = MultivariateNormalPrecision(jnp.zeros(3), Lambda)
>>> round(float(d.log_prob(jnp.ones(3))), 4)
-4.7171

Source code in src/gaussx/_distributions/_mvn_prec.py

class MultivariateNormalPrecision(AbstractMultivariateNormal):
    """Multivariate normal parameterized by a precision (inverse covariance) operator.

    This is the natural parameterization for many inference algorithms
    (e.g. message passing, variational inference in natural coordinates).
    The precision operator ``Lambda`` satisfies ``Lambda = Sigma^{-1}``.
    Precision-parameterised member of `AbstractMultivariateNormal`, which
    supplies the shared accessors (``covariance_operator``,
    ``precision_operator``, ``covariance_matrix``, ``scale_tril``, `kl`, ...).

    Requires the ``numpyro`` optional extra
    (``pip install "gaussx[numpyro]"``).

    Args:
        loc: Mean vector of shape ``(N,)``.
        prec_operator: Precision matrix as a lineax linear operator of
            shape ``(N, N)``.
        solver: Solver strategy for ``solve`` and ``logdet``. Defaults
            to ``AutoSolver()``.
        validate_args: Whether to validate input arguments.

    Examples:

        >>> import jax.numpy as jnp
        >>> import lineax as lx
        >>> from gaussx._distributions import MultivariateNormalPrecision
        >>> Lambda = lx.MatrixLinearOperator(
        ...     2.0 * jnp.eye(3), lx.positive_semidefinite_tag
        ... )
        >>> d = MultivariateNormalPrecision(jnp.zeros(3), Lambda)
        >>> round(float(d.log_prob(jnp.ones(3))), 4)
        -4.7171
    """

    pytree_data_fields = ("loc", "prec_operator", "solver")

    def __init__(
        self,
        loc: Float[Array, "*batch N"],
        prec_operator: lx.AbstractLinearOperator,
        solver: AbstractSolverStrategy | None = None,
        *,
        validate_args: bool | None = None,
    ) -> None:
        if solver is None:
            solver = AutoSolver()
        self.loc = loc
        self.prec_operator = prec_operator
        self.solver = solver
        event_shape = loc.shape[-1:]
        batch_shape = loc.shape[:-1]
        super().__init__(
            batch_shape=batch_shape,
            event_shape=event_shape,
            validate_args=validate_args,
        )

    @property
    def covariance_operator(self) -> lx.AbstractLinearOperator:
        """The covariance: the lazy `gaussx.inv` of ``prec_operator``."""
        return _inv(self.prec_operator)

    @property
    def precision_operator(self) -> lx.AbstractLinearOperator:
        """The precision operator (native; the same object as ``prec_operator``)."""
        return self.prec_operator

    def _log_prob_single(self, residual: Float[Array, " N"]) -> Float[Array, ""]:
        quad = jnp.sum(residual * self.prec_operator.mv(residual), axis=-1)
        ld = self.solver.logdet(self.prec_operator)
        n = self.loc.shape[-1]
        return -0.5 * (n * _LOG_2PI - ld + quad)

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*batch N"]:
        if key is None:
            raise ValueError(
                "PRNG key must be provided to sample from MultivariateNormalPrecision."
            )
        shape = sample_shape + self.batch_shape + self.event_shape
        # Draw in the parameters' dtype: default-float noise meets a float32
        # factor in the solve below, which lineax rejects under x64.
        dtype = jnp.result_type(self.loc, self.prec_operator.in_structure().dtype)
        eps = jax.random.normal(key, shape=shape, dtype=dtype)  # type: ignore[arg-type]
        eps_flat = rearrange(eps, "... D -> (...) D")

        # Look through tags and scalar multiples so a structured precision
        # (e.g. 2 * Kronecker) keeps its structured Cholesky factor.
        precision, scale = _unwrap_scaled(self.prec_operator)
        L = _cholesky(precision)

        def _solve_one(z):
            if isinstance(L, SparseCholeskyFactor):
                # A SparseOperator's factor is of P Λ Pᵀ: x = Pᵀ L⁻ᵀ z.
                return L.solve_lower_transpose(z)
            return _solve(L.T, z)

        samples_flat = scale * jax.vmap(_solve_one)(eps_flat)
        if isinstance(precision, lx.MatrixLinearOperator):
            # A dense Cholesky that fails numerically (an ill-conditioned or
            # only semi-definite precision) gives non-finite draws; take the
            # dense symmetric inverse square root then. Dense precisions only:
            # lax.cond stages both branches under jit, and a dense branch
            # would put N x N buffers into every structured or sparse sample.
            structured = samples_flat
            samples_flat = jax.lax.cond(
                jnp.all(jnp.isfinite(structured)),
                lambda: structured,
                lambda: _dense_precision_draws(self.prec_operator, eps_flat),
            )
        return self.loc + _reshape_samples(samples_flat, shape[:-1])

    @lazy_property
    def variance(self) -> Float[Array, "*batch N"]:
        # diag(Λ⁻¹) through the exact structured paths of `diag_inv` (e.g. the
        # O(N d³) selected inverse of a BlockTriDiag); only an operator with
        # none densifies (gh-360).
        variance = _diag_inv_structured(
            self.prec_operator, pinv=False, num_probes=30, key=None, solver=None
        )
        if variance is None:
            variance = _diag(_inv(self.prec_operator))
        return jnp.broadcast_to(variance, self.batch_shape + self.event_shape)

    def entropy(self) -> Float[Array, "*batch"]:
        n = self.loc.shape[-1]
        ld = self.solver.logdet(self.prec_operator)
        # One entry per batch element, as numpyro's contract requires.
        return jnp.broadcast_to(0.5 * (n * (1.0 + _LOG_2PI) - ld), self.batch_shape)

covariance_operator: lx.AbstractLinearOperator property

The covariance: the lazy gaussx.inv of prec_operator.

precision_operator: lx.AbstractLinearOperator property

The precision operator (native; the same object as prec_operator).

Gaussian Markov random fields

GaussianMRF is \(\mathcal N(\mu, Q^{-1})\) with a structured precision, and IntrinsicGMRF the improper \(\mathcal N(\mu, (\tau R)^+)\) of RW1 / RW2 / Besag (ICAR) priors, with a structure matrix \(R\) and its null space. Every operation dispatches on the precision's structure (GMRF precision builders return the right one):

• sample, in order: a factorisable precision (dense, BlockTriDiag, or a SparseOperator through its sparse Cholesky factor) gives \(\mu + L^{-\top}z\); a Kronecker-structured one (Kronecker, KroneckerSum, SpectralFunction) the symmetric root in its factor eigenvectors; anything else with precision_factors \(Q = \sum_k F_k^\top F_k\) is sampled by perturbation-optimisation: \(x = Q^{-1}\sum_k F_k^\top z_k\) has covariance \(Q^{-1}QQ^{-1} = Q^{-1}\), one CG solve per draw and no factorisation (graph precisions come factored, \(B^\top\operatorname{diag}(w)B\) from an incidence matrix).
• log_prob: \(\tfrac12\log|Q| - \tfrac12(x-\mu)^\top Q(x-\mu) - \tfrac N2\log 2\pi\), with \(\log|Q|\) from log_det_precision when known in closed form (no stochastic estimate inside an MCMC density), else logdet_strategy, solver or the structural logdet. The intrinsic density is \(\tfrac{N-c}{2}\log\tau - \tfrac\tau2(x-\mu)^\top R(x-\mu)\), plus the \(\tau\)-free \(\tfrac12\log|R|_+\) (pseudo_logdet) with include_normalizer.
• marginal variances: diag_inv (Takahashi on a sparse or banded factor, factor eigenvectors on a grid).
• observations \(y = Ax + \varepsilon\): condition_on_observations returns the posterior GMRF, \(Q + A^\top\Lambda A\) (sparse on the union pattern for a sparse prior and A, matrix-free for a grid prior) with its mean from one solve.
• constraints \(A_cx = e\): condition_on_constraints returns a ConstrainedGMRF, conditioning by kriging (Rue & Held, 2005, §2.3.3): \(x^\ast = x - Q^{-1}A_c^\top(A_cQ^{-1}A_c^\top)^{-1}(A_cx - e)\) is an exact draw from \(x \mid A_cx = e\), with corrected marginal variances and the constrained density. IntrinsicGMRF defaults to the hard constraint null_spaceᵀx = 0; constraint="soft" puts a tight Gaussian on null_spaceᵀx instead, for NUTS.

An odd-length rw2_structure(n) has one decoupled padding node: build the IntrinsicGMRF on all n + 1 nodes with null_space zero on it, and drop it from draws. All four require numpyro to be installed.

BYM2. Use BYM2GMRF(R_star, tau, phi, null_space) for the BYM2 pair \((b, u^\ast)\), not an IntrinsicGMRF on bym2_precision. The joint precision \(Q(\tau, \phi)\) has the null vector \((\sqrt{\phi/\tau}\,\mathbf 1, \mathbf 1)\), which depends on \(\theta\). So the bracket IntrinsicGMRF drops by default is not constant (it holds \(\tfrac n2\log\tfrac{\tau}{1-\phi}\)), and projecting draws along null_space does not remove the singular direction. BYM2GMRF uses the exact density on \(\{V^\top u^\ast = 0\}\), \(-\tfrac12 x^\top Q x + \tfrac n2\log\tfrac{\tau}{1-\phi}\) (plus the constant \(\tfrac12\log|R^\ast|_+ - \tfrac{2n-c}{2}\log 2\pi\) with include_normalizer), and draws \(u^\ast\) from the ICAR on \(R^\ast\), then \(b\) from its definition. It is an IntrinsicGMRF, so laplace_mode takes it as a prior, and its log-marginal is then correct in \((\tau, \phi)\):

scale = gx.generalized_variance_scale(R, V)  # R: the Besag structure, n areas
R_star = scale * R
# theta-free, so once: log|s R|_+ = log|R|_+ + (n - c) log s, with c = 1 here
log_pdet = gx.pseudo_logdet(R, structure="laplacian") + (n - 1) * jnp.log(scale)


def log_marginal(tau, phi):
    prior = gx.BYM2GMRF(R_star, tau, phi, V, include_normalizer=True, log_pdet=log_pdet)
    return gx.laplace_mode(
        prior, gx.PoissonLikelihood(counts), projector=A
    ).log_marginal

import einx
import jax
import jax.numpy as jnp
import numpy as np

import gaussx as gx

# Gap-fill a cloudy sea-surface-temperature snapshot on an H × W grid
H, W = 30, 40
kappa, tau, alpha = gx.matern_spde_params(range=8.0, sigma=1.0, nu=1.0, d=2)
prior = gx.GaussianMRF(
    jnp.zeros(H * W), gx.spde_precision_grid((H, W), kappa, tau, alpha=alpha)
)
clear_idx = np.flatnonzero(np.random.default_rng(0).random(H * W) > 0.6)
n_clear = clear_idx.size
sst = jnp.sin(jnp.arange(H * W) / 50.0)  # stand-in for the measured field
A = gx.SparseOperator.from_coo(
    np.arange(n_clear), clear_idx, jnp.ones(n_clear), (n_clear, H * W)
)
post = prior.condition_on_observations(A, noise_precision=1 / 0.2**2, y=sst[clear_idx])
filled = einx.id("(h w) -> h w", post.loc, h=H)
sd = einx.id("(h w) -> h w", jnp.sqrt(post.marginal_variances()), h=H)
draws = post.sample(jax.random.PRNGKey(0), (20,))  # 20 plausible gap-filled fields

# An ICAR field with a hard sum-to-zero constraint per island group: the
# mainland path 0 - 1 - 2 - 3 and the island pair 4 - 5
senders, receivers = np.array([1, 2, 3, 5]), np.array([0, 1, 2, 4])
degree = np.bincount(np.r_[senders, receivers], minlength=6).astype(float)
R = gx.SparseOperator.from_coo(
    np.r_[np.arange(6), senders],
    np.r_[np.arange(6), receivers],
    jnp.asarray(np.r_[degree, -np.ones(4)]),
    (6, 6),
    symmetric=True,
)
null_space = jnp.zeros((6, 2)).at[:4, 0].set(0.5).at[4:, 1].set(2**-0.5)
icar = gx.IntrinsicGMRF(jnp.zeros(6), 2.0, R, null_space)  # kernellib: graph_null_space
u = icar.sample(jax.random.PRNGKey(1))  # null_spaceᵀ u == 0 to machine precision
lp = icar.log_prob(u)  # (N − c)/2 · log τ − τ/2 uᵀRu, with N − c = 4

Structured linear algebra and Gaussian primitives for JAX.

GaussianMRF

Bases: Distribution

Gaussian Markov random field \(x \sim \mathcal N(\mu, Q^{-1})\).

The precision is a structured lineax operator, and log_prob, sample and marginal_variances take its cheapest exact path:

• SparseOperator: sparse Cholesky, μ + Pᵀ L⁻ᵀ z; Takahashi for the marginal variances.
• BlockTriDiag: block Cholesky and the block selected inverse.
• Kronecker, KroneckerSum, SpectralFunction: the factor eigenvectors, V⁻¹ Λ^{-1/2} V z, and the factor eigenvalues.
• dense: Cholesky.
• anything else, with precision_factors: perturbation-optimisation, one solve of Q x = Σ_k F_kᵀ z_k per draw; solver / logdet_strategy for log|Q|.

A SparseOperator defaults to SparseCholeskySolver for every solve and log-determinant; pass precision_factors (and no sparse Cholesky solver) to sample matrix-free by CG instead (solver, if given, does that solve). Graph precisions come factored — Bᵀ diag(w) B from an incidence matrix, proper CAR, Leroux — so a 10⁶-node field samples with no factorisation at all.

Requires the numpyro optional extra. loc is a single field (batch_shape == ()); log_prob takes any leading sample axes.

Parameters:

Name	Type	Description	Default
loc	Float[ArrayLike, ' N']	Mean \(\mu\), shape (N,).	required
precision	AbstractLinearOperator	\(Q\), a symmetric positive-definite operator of size N.	required
precision_factors	tuple[AbstractLinearOperator, ...] | None	Operators \(F_k\) with \(Q = \sum_k F_k^\top F_k\), for the perturbation-optimisation sampler.	None
solver	AbstractSolverStrategy | None	Solve (and log-determinant) strategy, e.g. SparseCholeskySolver or CGSolver.	None
logdet_strategy	AbstractLogdetStrategy | None	Log-determinant strategy; takes precedence over solver for log|Q| (e.g. SLQLogdet on a huge field).	None
log_det_precision	Float[ArrayLike, ''] | None	A known \(\log|Q|\), which overrides both (callers with a closed form, such as a proper CAR, skip a stochastic estimate inside an MCMC log-density).	None
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

import jax
import jax.numpy as jnp
import numpy as np
import gaussx as gx

# A proper field on the path graph 0 - 1 - 2 - 3 - 4: Q = R + I
n = 5
Q = gx.SparseOperator.from_coo(
    np.r_[np.arange(n), np.arange(1, n)],
    np.r_[np.arange(n), np.arange(n - 1)],
    jnp.r_[jnp.array([2.0, 3.0, 3.0, 3.0, 2.0]), -jnp.ones(n - 1)],
    (n, n),
    symmetric=True,
)
prior = gx.GaussianMRF(jnp.zeros(n), Q)  # sparse Cholesky throughout
x = prior.sample(jax.random.PRNGKey(0), (3,))  # (3, 5)
lp = prior.log_prob(x)  # (3,)
sd = jnp.sqrt(prior.marginal_variances())  # Takahashi

# Observe nodes 0 and 3 with noise sd 0.1: the posterior is a GMRF
A = gx.SparseOperator.from_coo([0, 1], [0, 3], jnp.ones(2), (2, n))
post = prior.condition_on_observations(A, 1 / 0.1**2, jnp.array([1.0, -1.0]))

# A sum-to-zero field, by kriging
zero_sum = prior.condition_on_constraints(jnp.ones((1, n)), jnp.zeros(1))
u = zero_sum.sample(jax.random.PRNGKey(1))  # u.sum() == 0

Source code in src/gaussx/_distributions/_gmrf.py

class GaussianMRF(dist.Distribution):
    r"""Gaussian Markov random field $x \sim \mathcal N(\mu, Q^{-1})$.

    The precision is a structured lineax operator, and ``log_prob``,
    ``sample`` and ``marginal_variances`` take its cheapest exact path:

    - `SparseOperator`: sparse Cholesky, ``μ + Pᵀ L⁻ᵀ z``; Takahashi for the
      marginal variances.
    - `BlockTriDiag`: block Cholesky and the block selected inverse.
    - `Kronecker`, `KroneckerSum`, `SpectralFunction`: the factor
      eigenvectors, ``V⁻¹ Λ^{-1/2} V z``, and the factor eigenvalues.
    - dense: Cholesky.
    - anything else, with ``precision_factors``: perturbation-optimisation,
      one solve of ``Q x = Σ_k F_kᵀ z_k`` per draw; ``solver`` /
      ``logdet_strategy`` for ``log|Q|``.

    A `SparseOperator` defaults to `SparseCholeskySolver` for every solve and
    log-determinant; pass ``precision_factors`` (and no sparse Cholesky
    solver) to sample matrix-free by CG instead (``solver``, if given, does
    that solve). Graph precisions come factored — ``Bᵀ diag(w) B`` from an
    incidence matrix, proper CAR, Leroux — so a 10⁶-node field samples with
    no factorisation at all.

    Requires the ``numpyro`` optional extra. ``loc`` is a single field
    (``batch_shape == ()``); ``log_prob`` takes any leading sample axes.

    Args:
        loc: Mean $\mu$, shape ``(N,)``.
        precision: $Q$, a symmetric positive-definite operator of size ``N``.
        precision_factors: Operators $F_k$ with $Q = \sum_k F_k^\top F_k$,
            for the perturbation-optimisation sampler.
        solver: Solve (and log-determinant) strategy, e.g.
            `SparseCholeskySolver` or `CGSolver`.
        logdet_strategy: Log-determinant strategy; takes precedence over
            ``solver`` for ``log|Q|`` (e.g. `SLQLogdet` on a huge field).
        log_det_precision: A known $\log|Q|$, which overrides both (callers
            with a closed form, such as a proper CAR, skip a stochastic
            estimate inside an MCMC log-density).
        validate_args: Whether to validate input arguments.

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

        # A proper field on the path graph 0 - 1 - 2 - 3 - 4: Q = R + I
        n = 5
        Q = gx.SparseOperator.from_coo(
            np.r_[np.arange(n), np.arange(1, n)],
            np.r_[np.arange(n), np.arange(n - 1)],
            jnp.r_[jnp.array([2.0, 3.0, 3.0, 3.0, 2.0]), -jnp.ones(n - 1)],
            (n, n),
            symmetric=True,
        )
        prior = gx.GaussianMRF(jnp.zeros(n), Q)  # sparse Cholesky throughout
        x = prior.sample(jax.random.PRNGKey(0), (3,))  # (3, 5)
        lp = prior.log_prob(x)  # (3,)
        sd = jnp.sqrt(prior.marginal_variances())  # Takahashi

        # Observe nodes 0 and 3 with noise sd 0.1: the posterior is a GMRF
        A = gx.SparseOperator.from_coo([0, 1], [0, 3], jnp.ones(2), (2, n))
        post = prior.condition_on_observations(A, 1 / 0.1**2, jnp.array([1.0, -1.0]))

        # A sum-to-zero field, by kriging
        zero_sum = prior.condition_on_constraints(jnp.ones((1, n)), jnp.zeros(1))
        u = zero_sum.sample(jax.random.PRNGKey(1))  # u.sum() == 0
        ```
    """

    arg_constraints = {"loc": dist.constraints.real_vector}  # noqa: RUF012
    support = dist.constraints.real_vector
    reparametrized_params = ["loc"]  # noqa: RUF012
    pytree_data_fields = (
        "loc",
        "precision",
        "precision_factors",
        "solver",
        "logdet_strategy",
        "log_det_precision",
    )

    def __init__(
        self,
        loc: Float[ArrayLike, " N"],
        precision: lx.AbstractLinearOperator,
        precision_factors: tuple[lx.AbstractLinearOperator, ...] | None = None,
        solver: AbstractSolverStrategy | None = None,
        logdet_strategy: AbstractLogdetStrategy | None = None,
        log_det_precision: Float[ArrayLike, ""] | None = None,
        *,
        validate_args: bool | None = None,
    ) -> None:
        loc = jnp.asarray(loc)
        n = precision.in_size()
        if loc.shape != (n,):
            raise ValueError(
                f"loc must have shape ({n},) to match the precision, got {loc.shape}."
            )
        self.loc = loc
        self.precision = precision
        self.precision_factors = (
            None if precision_factors is None else tuple(precision_factors)
        )
        self.solver = solver
        self.logdet_strategy = logdet_strategy
        self.log_det_precision = (
            None if log_det_precision is None else jnp.asarray(log_det_precision)
        )
        super().__init__(batch_shape=(), event_shape=(n,), validate_args=validate_args)

    # -- solves ---------------------------------------------------------------

    def _log_det(self) -> Float[Array, ""]:
        """``log|Q|``: ``log_det_precision``, else the strategy, else `logdet`.

        Returns:
            Scalar log-determinant of the precision.
        """
        if self.log_det_precision is not None:
            return self.log_det_precision
        if self.logdet_strategy is not None:
            return self.logdet_strategy.logdet(self.precision)
        return dispatch_logdet(
            self.precision, _solve_strategy(self.precision, self.solver)
        )

    # -- density --------------------------------------------------------------

    @validate_sample
    def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
        r"""$\tfrac12\log|Q| - \tfrac12(x-\mu)^\top Q(x-\mu) - \tfrac N2\log 2\pi$.

        Args:
            value: Points, shape ``(..., N)``.

        Returns:
            Log-densities, shape ``(...)``.
        """
        n = self.event_shape[0]
        quad = _quadratic(self.precision, value - self.loc)
        return 0.5 * (self._log_det() - quad - n * _LOG_2PI)

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*sample N"]:
        """Exact draws through the structural sampling dispatch.

        Args:
            key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
            sample_shape: Leading sample shape.

        Returns:
            Draws, shape ``sample_shape + (N,)``.
        """
        size, draw = _precision_sampler(
            self.precision, self.solver, self.precision_factors
        )
        draws = _draw_samples(key, sample_shape, size, draw, self.loc.dtype)
        return self.loc + draws

    def marginal_variances(self) -> Float[Array, " N"]:
        """``diag(Q⁻¹)`` through `diag_inv`'s structural dispatch.

        Returns:
            Marginal variances, shape ``(N,)``.
        """
        strategy = _solve_strategy(self.precision, self.solver)
        return diag_inv(self.precision, solver=strategy)

    @lazy_property
    def mean(self) -> Float[Array, " N"]:
        return self.loc

    @lazy_property
    def variance(self) -> Float[Array, " N"]:
        return self.marginal_variances()

    # -- conditioning ---------------------------------------------------------

    def condition_on_observations(
        self,
        A: lx.AbstractLinearOperator | Float[ArrayLike, "M N"],
        noise_precision: Float[ArrayLike, ""] | Float[ArrayLike, " M"],
        y: Float[ArrayLike, " M"],
    ) -> GaussianMRF:
        r"""The posterior given $y = Ax + \varepsilon$ with noise precision $\Lambda$.

        $Q_\text{post} = Q + A^\top\Lambda A$ and
        $\mu_\text{post} = \mu + Q_\text{post}^{-1}A^\top\Lambda(y - A\mu)$,
        one solve. The posterior precision keeps what structure it can:

        - a `SparseOperator` prior observed through a `SparseOperator` ``A``
          stays sparse, on the union of the two patterns
          (`SparseOperator.congruence`, planned once per pair of patterns),
          so it keeps the sparse Cholesky path;
        - a dense prior stays dense;
        - anything else (a grid `SpectralFunction`, a `KroneckerSum`, a
          `BlockTriDiag` observed at scattered points) becomes the
          matrix-free sum ``Q + (Λ^{1/2}A)ᵀ(Λ^{1/2}A)``: solves by `solve`
          (CG when large), and draws by perturbation-optimisation with the
          factors ``(Q^{1/2}, Λ^{1/2}A)``, ``Q^{1/2}`` the symmetric root of
          a `SpectralFunction` or `KroneckerSum` prior.

        ``precision_factors`` gain ``Λ^{1/2}A``; a known ``log_det_precision``
        is dropped.

        Args:
            A: Observation operator, shape ``(M, N)`` (array or operator).
            noise_precision: $\Lambda$'s diagonal, a scalar or shape ``(M,)``.
            y: Observations, shape ``(M,)``.

        Returns:
            The posterior `GaussianMRF`, with the same solver strategies (a
            `SparseCholeskySolver` is dropped if the posterior is not sparse).
        """
        A = _as_operator(A)
        m = A.out_size()
        dtype = self.loc.dtype
        w = jnp.broadcast_to(jnp.asarray(noise_precision, dtype=dtype), (m,))
        root_A = lx.DiagonalLinearOperator(jnp.sqrt(w)) @ A
        prior = _unwrap(self.precision)
        factors = self.precision_factors
        if isinstance(prior, SparseOperator) and isinstance(A, SparseOperator):
            q_post: lx.AbstractLinearOperator = prior.union(
                prior.congruence(A, w), tags=_PSD
            )
        elif isinstance(prior, lx.MatrixLinearOperator):
            dense = A.as_matrix()
            weighted = einx.multiply("k i, k -> k i", dense, w)
            update = einsum(weighted, dense, "k i, k j -> i j")
            q_post = lx.MatrixLinearOperator(prior.matrix + update, _PSD)
        else:
            q_post = lx.TaggedLinearOperator(
                lx.AddLinearOperator(self.precision, root_A.T @ root_A), _PSD
            )
            if factors is None:
                root = _symmetric_root(prior)
                factors = None if root is None else (root,)
        if factors is not None:
            factors = (*factors, root_A)
        solver = self.solver
        if isinstance(solver, SparseCholeskySolver) and not isinstance(
            q_post, SparseOperator
        ):
            solver = None
        residual = jnp.asarray(y, dtype=dtype) - A.mv(self.loc)
        rhs = A.transpose().mv(w * residual)
        loc = self.loc + _solve_fn(q_post, solver)(rhs)
        return GaussianMRF(
            loc, q_post, factors, solver=solver, logdet_strategy=self.logdet_strategy
        )

    def condition_on_constraints(
        self,
        A_c: Float[ArrayLike, "c N"] | Float[ArrayLike, " N"],
        e: Float[ArrayLike, " c"] | Float[ArrayLike, ""],
    ) -> ConstrainedGMRF:
        """Hard linear constraints ``A_c x = e``, by conditioning by kriging.

        Args:
            A_c: Constraint matrix, shape ``(c, N)`` (or ``(N,)`` for one).
            e: Constraint values, shape ``(c,)``.

        Returns:
            The `ConstrainedGMRF` ``x | A_c x = e``.
        """
        return ConstrainedGMRF(self, A_c, e)

log_prob(value: Float[Array, '*sample N']) -> Float[Array, '*sample']

\(\tfrac12\log|Q| - \tfrac12(x-\mu)^\top Q(x-\mu) - \tfrac N2\log 2\pi\).

Parameters:

Name	Type	Description	Default
value	Float[Array, '*sample N']	Points, shape (..., N).	required

Returns:

Type	Description
Float[Array, '*sample']	Log-densities, shape (...).

Source code in src/gaussx/_distributions/_gmrf.py

@validate_sample
def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
    r"""$\tfrac12\log|Q| - \tfrac12(x-\mu)^\top Q(x-\mu) - \tfrac N2\log 2\pi$.

    Args:
        value: Points, shape ``(..., N)``.

    Returns:
        Log-densities, shape ``(...)``.
    """
    n = self.event_shape[0]
    quad = _quadratic(self.precision, value - self.loc)
    return 0.5 * (self._log_det() - quad - n * _LOG_2PI)

sample(key: jax.Array | None, sample_shape: tuple[int, ...] = ()) -> Float[Array, '*sample N']

Exact draws through the structural sampling dispatch.

Parameters:

Name	Type	Description	Default
key	Array | None	PRNG key. None means jax.random.PRNGKey(0).	required
sample_shape	tuple[int, ...]	Leading sample shape.	()

Returns:

Type	Description
Float[Array, '*sample N']	Draws, shape sample_shape + (N,).

Source code in src/gaussx/_distributions/_gmrf.py

def sample(
    self,
    key: jax.Array | None,
    sample_shape: tuple[int, ...] = (),
) -> Float[Array, "*sample N"]:
    """Exact draws through the structural sampling dispatch.

    Args:
        key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
        sample_shape: Leading sample shape.

    Returns:
        Draws, shape ``sample_shape + (N,)``.
    """
    size, draw = _precision_sampler(
        self.precision, self.solver, self.precision_factors
    )
    draws = _draw_samples(key, sample_shape, size, draw, self.loc.dtype)
    return self.loc + draws

marginal_variances() -> Float[Array, ' N']

diag(Q⁻¹) through diag_inv's structural dispatch.

Returns:

Type	Description
Float[Array, ' N']	Marginal variances, shape (N,).

Source code in src/gaussx/_distributions/_gmrf.py

def marginal_variances(self) -> Float[Array, " N"]:
    """``diag(Q⁻¹)`` through `diag_inv`'s structural dispatch.

    Returns:
        Marginal variances, shape ``(N,)``.
    """
    strategy = _solve_strategy(self.precision, self.solver)
    return diag_inv(self.precision, solver=strategy)

condition_on_observations(A: lx.AbstractLinearOperator | Float[ArrayLike, 'M N'], noise_precision: Float[ArrayLike, ''] | Float[ArrayLike, ' M'], y: Float[ArrayLike, ' M']) -> GaussianMRF

The posterior given \(y = Ax + \varepsilon\) with noise precision \(\Lambda\).

\(Q_\text{post} = Q + A^\top\Lambda A\) and \(\mu_\text{post} = \mu + Q_\text{post}^{-1}A^\top\Lambda(y - A\mu)\), one solve. The posterior precision keeps what structure it can:

• a SparseOperator prior observed through a SparseOperator A stays sparse, on the union of the two patterns (SparseOperator.congruence, planned once per pair of patterns), so it keeps the sparse Cholesky path;
• a dense prior stays dense;
• anything else (a grid SpectralFunction, a KroneckerSum, a BlockTriDiag observed at scattered points) becomes the matrix-free sum Q + (Λ^{1/2}A)ᵀ(Λ^{1/2}A): solves by solve (CG when large), and draws by perturbation-optimisation with the factors (Q^{1/2}, Λ^{1/2}A), Q^{1/2} the symmetric root of a SpectralFunction or KroneckerSum prior.

precision_factors gain Λ^{1/2}A; a known log_det_precision is dropped.

Parameters:

Name	Type	Description	Default
A	AbstractLinearOperator | Float[ArrayLike, 'M N']	Observation operator, shape (M, N) (array or operator).	required
noise_precision	Float[ArrayLike, ''] | Float[ArrayLike, ' M']	\(\Lambda\)'s diagonal, a scalar or shape (M,).	required
y	Float[ArrayLike, ' M']	Observations, shape (M,).	required

Returns:

Type	Description
GaussianMRF	The posterior GaussianMRF, with the same solver strategies (a
GaussianMRF	SparseCholeskySolver is dropped if the posterior is not sparse).

Source code in src/gaussx/_distributions/_gmrf.py

def condition_on_observations(
    self,
    A: lx.AbstractLinearOperator | Float[ArrayLike, "M N"],
    noise_precision: Float[ArrayLike, ""] | Float[ArrayLike, " M"],
    y: Float[ArrayLike, " M"],
) -> GaussianMRF:
    r"""The posterior given $y = Ax + \varepsilon$ with noise precision $\Lambda$.

    $Q_\text{post} = Q + A^\top\Lambda A$ and
    $\mu_\text{post} = \mu + Q_\text{post}^{-1}A^\top\Lambda(y - A\mu)$,
    one solve. The posterior precision keeps what structure it can:

    - a `SparseOperator` prior observed through a `SparseOperator` ``A``
      stays sparse, on the union of the two patterns
      (`SparseOperator.congruence`, planned once per pair of patterns),
      so it keeps the sparse Cholesky path;
    - a dense prior stays dense;
    - anything else (a grid `SpectralFunction`, a `KroneckerSum`, a
      `BlockTriDiag` observed at scattered points) becomes the
      matrix-free sum ``Q + (Λ^{1/2}A)ᵀ(Λ^{1/2}A)``: solves by `solve`
      (CG when large), and draws by perturbation-optimisation with the
      factors ``(Q^{1/2}, Λ^{1/2}A)``, ``Q^{1/2}`` the symmetric root of
      a `SpectralFunction` or `KroneckerSum` prior.

    ``precision_factors`` gain ``Λ^{1/2}A``; a known ``log_det_precision``
    is dropped.

    Args:
        A: Observation operator, shape ``(M, N)`` (array or operator).
        noise_precision: $\Lambda$'s diagonal, a scalar or shape ``(M,)``.
        y: Observations, shape ``(M,)``.

    Returns:
        The posterior `GaussianMRF`, with the same solver strategies (a
        `SparseCholeskySolver` is dropped if the posterior is not sparse).
    """
    A = _as_operator(A)
    m = A.out_size()
    dtype = self.loc.dtype
    w = jnp.broadcast_to(jnp.asarray(noise_precision, dtype=dtype), (m,))
    root_A = lx.DiagonalLinearOperator(jnp.sqrt(w)) @ A
    prior = _unwrap(self.precision)
    factors = self.precision_factors
    if isinstance(prior, SparseOperator) and isinstance(A, SparseOperator):
        q_post: lx.AbstractLinearOperator = prior.union(
            prior.congruence(A, w), tags=_PSD
        )
    elif isinstance(prior, lx.MatrixLinearOperator):
        dense = A.as_matrix()
        weighted = einx.multiply("k i, k -> k i", dense, w)
        update = einsum(weighted, dense, "k i, k j -> i j")
        q_post = lx.MatrixLinearOperator(prior.matrix + update, _PSD)
    else:
        q_post = lx.TaggedLinearOperator(
            lx.AddLinearOperator(self.precision, root_A.T @ root_A), _PSD
        )
        if factors is None:
            root = _symmetric_root(prior)
            factors = None if root is None else (root,)
    if factors is not None:
        factors = (*factors, root_A)
    solver = self.solver
    if isinstance(solver, SparseCholeskySolver) and not isinstance(
        q_post, SparseOperator
    ):
        solver = None
    residual = jnp.asarray(y, dtype=dtype) - A.mv(self.loc)
    rhs = A.transpose().mv(w * residual)
    loc = self.loc + _solve_fn(q_post, solver)(rhs)
    return GaussianMRF(
        loc, q_post, factors, solver=solver, logdet_strategy=self.logdet_strategy
    )

condition_on_constraints(A_c: Float[ArrayLike, 'c N'] | Float[ArrayLike, ' N'], e: Float[ArrayLike, ' c'] | Float[ArrayLike, '']) -> ConstrainedGMRF

Hard linear constraints A_c x = e, by conditioning by kriging.

Parameters:

Name	Type	Description	Default
A_c	Float[ArrayLike, 'c N'] | Float[ArrayLike, ' N']	Constraint matrix, shape (c, N) (or (N,) for one).	required
e	Float[ArrayLike, ' c'] | Float[ArrayLike, '']	Constraint values, shape (c,).	required

Returns:

Type	Description
ConstrainedGMRF	The ConstrainedGMRF x | A_c x = e.

Source code in src/gaussx/_distributions/_gmrf.py

def condition_on_constraints(
    self,
    A_c: Float[ArrayLike, "c N"] | Float[ArrayLike, " N"],
    e: Float[ArrayLike, " c"] | Float[ArrayLike, ""],
) -> ConstrainedGMRF:
    """Hard linear constraints ``A_c x = e``, by conditioning by kriging.

    Args:
        A_c: Constraint matrix, shape ``(c, N)`` (or ``(N,)`` for one).
        e: Constraint values, shape ``(c,)``.

    Returns:
        The `ConstrainedGMRF` ``x | A_c x = e``.
    """
    return ConstrainedGMRF(self, A_c, e)

IntrinsicGMRF

Bases: Distribution

Improper \(\mathcal N(\mu, (\tau R)^+)\) on the complement of \(\ker R\).

\(R\) is a positive-semidefinite structure matrix (RW1, RW2, Besag / ICAR, a grid Laplacian) with null space spanned by the columns of null_space (c of them), and \(\tau\) is the precision scale. The log-density is

\[ \log\pi(x) = \tfrac{N-c}{2}\log\tau - \tfrac\tau2(x-\mu)^\top R(x-\mu) \;\Big[+\ \tfrac12\log|R|_+ - \tfrac{N-c}{2}\log 2\pi\Big], \]

the bracket only with include_normalizer (it is \(\tau\)-free, so an MCMC target can skip it; pseudo_logdet computes it per call). The \(\tau\) term is always included, which is what makes \(\tau\) identifiable under a hyperprior. N − c must be rank(R). The bracket is constant only when \(R\) itself does not depend on the hyperparameters: for BYM2, whose structure holds \((\tau, \phi)\), use BYM2GMRF.

constraint chooses what happens on \(\ker R\):

• "hard" (default, what INLA needs): the field lives on \(\{V^\top x = 0\}\), \(V\) = null_space. The log-density above (with the normaliser) is then exactly the constrained density on that subspace. Samples are kriged onto it: R-INLA's recipe of factoring \(\tau R + \varepsilon I\) (with \(\varepsilon = \sqrt{\text{eps}}\) times the mean diagonal) and kriging against \(V^\top x = 0\) reduces, since \((\tau R + \varepsilon I)^{-1}V = V/\varepsilon\), to the orthogonal projection \(x - V(V^\top V)^{-1}V^\top x\), which is what is applied (exactly, with no extra solves). Marginal variances are kriging- corrected (ConstrainedGMRF) on the shifted factor, or exact from the pseudo-inverse on a Kronecker-structured \(R\).
• "soft" (for NUTS, which cannot sample a hard constraint): adds a tight Gaussian \(V^\top x \sim \mathcal N(0, s^2 I)\), \(s\) = soft_constraint_scale, to the density and to the samples.
• "none": only $\mu + $ a draw from \(\mathcal N(0, (\tau R)^+)\) — the density is improper and invariant along \(\ker R\).

The draw from \(\mathcal N(0, (\tau R)^+)\) dispatches like GaussianMRF: Cholesky of \(\tau R + \varepsilon I\) (sparse, banded or dense), the pseudo-inverse root in the factored eigenbasis (exact, no shift), or perturbation-optimisation with precision_factors (\(R = \sum_k F_k^\top F_k\), so the right-hand side lies in \(\operatorname{range}(R)\) and CG converges on the singular system); each is then projected orthogonally to null_space.

Odd-n RW2. rw2_structure with odd n returns n + 1 rows, the last a decoupled unit-precision node. Build the field on all N = n + 1 nodes with null_space zero on the padding row (c = 2): then N − c = n − 1 = rank(R), and the padding node is an independent \(\mathcal N(0, 1/\tau)\) coordinate whose own ½ log τ is the one extra in (N − c)/2 · log τ. Sample it with the rest and drop it (x[:n]).

Requires the numpyro optional extra.

Parameters:

Name	Type	Description	Default
loc	Float[ArrayLike, ' N']	\(\mu\), shape (N,).	required
precision_scale	Float[ArrayLike, '']	\(\tau > 0\).	required
structure	AbstractLinearOperator	\(R\), symmetric positive semidefinite, size N.	required
null_space	Float[ArrayLike, 'N c'] | Float[ArrayLike, ' N']	A basis of \(\ker R\), shape (N, c) (or (N,)); for a graph, one column per connected component (kernellib's graph_null_space).	required
precision_factors	tuple[AbstractLinearOperator, ...] | None	\(F_k\) with \(R = \sum_k F_k^\top F_k\) (e.g. the weighted incidence matrix), for the CG sampler.	None
constraint	Literal['none', 'soft', 'hard']	"hard", "soft" or "none".	'hard'
soft_constraint_scale	float	\(s\) for "soft".	0.001
include_normalizer	bool	Add the \(\tau\)-free constant (exact for an orthonormal null_space under "soft").	False
validate_args	bool | None	Whether to validate input arguments.	None

Raises:

Type	Description
ValueError	For an unknown constraint or mismatched shapes.

Examples:

import jax
import jax.numpy as jnp
import gaussx as gx

# An RW1 trend with a hard sum-to-zero constraint
n = 50
R = gx.rw1_structure(n)
trend = gx.IntrinsicGMRF(jnp.zeros(n), 4.0, R, jnp.ones((n, 1)) / jnp.sqrt(n))
u = trend.sample(jax.random.PRNGKey(0), (5,))  # rows sum to 0
lp = trend.log_prob(u)  # includes (n - 1)/2 · log 4

# RW2 with odd n: one padding node, null space {1, t} zero on it
R2 = gx.rw2_structure(7)  # 8 x 8
t = jnp.arange(7.0)
V = jnp.zeros((8, 2)).at[:7, 0].set(1.0).at[:7, 1].set(t - t.mean())
smooth = gx.IntrinsicGMRF(jnp.zeros(8), 10.0, R2, V)
x = smooth.sample(jax.random.PRNGKey(1))[:7]  # drop the padding node

Source code in src/gaussx/_distributions/_gmrf.py

class IntrinsicGMRF(dist.Distribution):
    r"""Improper $\mathcal N(\mu, (\tau R)^+)$ on the complement of $\ker R$.

    $R$ is a positive-semidefinite structure matrix (RW1, RW2, Besag / ICAR,
    a grid Laplacian) with null space spanned by the columns of
    ``null_space`` (``c`` of them), and $\tau$ is the precision scale. The
    log-density is

    $$
    \log\pi(x) = \tfrac{N-c}{2}\log\tau - \tfrac\tau2(x-\mu)^\top R(x-\mu)
        \;\Big[+\ \tfrac12\log|R|_+ - \tfrac{N-c}{2}\log 2\pi\Big],
    $$

    the bracket only with ``include_normalizer`` (it is $\tau$-free, so an
    MCMC target can skip it; `pseudo_logdet` computes it per call). The
    $\tau$ term is always included, which is what makes $\tau$ identifiable
    under a hyperprior. ``N − c`` must be ``rank(R)``. The bracket is
    constant only when $R$ itself does not depend on the hyperparameters:
    for BYM2, whose structure holds $(\tau, \phi)$, use `BYM2GMRF`.

    ``constraint`` chooses what happens on $\ker R$:

    - ``"hard"`` (default, what INLA needs): the field lives on
      $\{V^\top x = 0\}$, $V$ = ``null_space``. The log-density above (with
      the normaliser) is then exactly the constrained density on that
      subspace. Samples are kriged onto it: R-INLA's recipe of factoring
      $\tau R + \varepsilon I$ (with $\varepsilon = \sqrt{\text{eps}}$ times
      the mean diagonal) and kriging against $V^\top x = 0$ reduces, since
      $(\tau R + \varepsilon I)^{-1}V = V/\varepsilon$, to the orthogonal
      projection $x - V(V^\top V)^{-1}V^\top x$, which is what is applied
      (exactly, with no extra solves). Marginal variances are kriging-
      corrected (`ConstrainedGMRF`) on the shifted factor, or exact from
      the pseudo-inverse on a Kronecker-structured $R$.
    - ``"soft"`` (for NUTS, which cannot sample a hard constraint): adds a
      tight Gaussian $V^\top x \sim \mathcal N(0, s^2 I)$,
      $s$ = ``soft_constraint_scale``, to the density and to the samples.
    - ``"none"``: only $\mu + $ a draw from $\mathcal N(0, (\tau R)^+)$ —
      the density is improper and invariant along $\ker R$.

    The draw from $\mathcal N(0, (\tau R)^+)$ dispatches like `GaussianMRF`:
    Cholesky of $\tau R + \varepsilon I$ (sparse, banded or dense), the
    pseudo-inverse root in the factored eigenbasis (exact, no shift), or
    perturbation-optimisation with ``precision_factors``
    ($R = \sum_k F_k^\top F_k$, so the right-hand side lies in
    $\operatorname{range}(R)$ and CG converges on the singular system); each
    is then projected orthogonally to ``null_space``.

    **Odd-``n`` RW2.** `rw2_structure` with odd ``n`` returns ``n + 1`` rows,
    the last a decoupled unit-precision node. Build the field on all
    ``N = n + 1`` nodes with ``null_space`` zero on the padding row
    (``c = 2``): then ``N − c = n − 1 = rank(R)``, and the padding node is an
    independent $\mathcal N(0, 1/\tau)$ coordinate whose own ½ log τ is the
    one extra in ``(N − c)/2 · log τ``. Sample it with the rest and drop it
    (``x[:n]``).

    Requires the ``numpyro`` optional extra.

    Args:
        loc: $\mu$, shape ``(N,)``.
        precision_scale: $\tau > 0$.
        structure: $R$, symmetric positive semidefinite, size ``N``.
        null_space: A basis of $\ker R$, shape ``(N, c)`` (or ``(N,)``); for a
            graph, one column per connected component (kernellib's
            ``graph_null_space``).
        precision_factors: $F_k$ with $R = \sum_k F_k^\top F_k$ (e.g. the
            weighted incidence matrix), for the CG sampler.
        constraint: ``"hard"``, ``"soft"`` or ``"none"``.
        soft_constraint_scale: $s$ for ``"soft"``.
        include_normalizer: Add the $\tau$-free constant (exact for an
            orthonormal ``null_space`` under ``"soft"``).
        validate_args: Whether to validate input arguments.

    Raises:
        ValueError: For an unknown ``constraint`` or mismatched shapes.

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

        # An RW1 trend with a hard sum-to-zero constraint
        n = 50
        R = gx.rw1_structure(n)
        trend = gx.IntrinsicGMRF(jnp.zeros(n), 4.0, R, jnp.ones((n, 1)) / jnp.sqrt(n))
        u = trend.sample(jax.random.PRNGKey(0), (5,))  # rows sum to 0
        lp = trend.log_prob(u)  # includes (n - 1)/2 · log 4

        # RW2 with odd n: one padding node, null space {1, t} zero on it
        R2 = gx.rw2_structure(7)  # 8 x 8
        t = jnp.arange(7.0)
        V = jnp.zeros((8, 2)).at[:7, 0].set(1.0).at[:7, 1].set(t - t.mean())
        smooth = gx.IntrinsicGMRF(jnp.zeros(8), 10.0, R2, V)
        x = smooth.sample(jax.random.PRNGKey(1))[:7]  # drop the padding node
        ```
    """

    arg_constraints = {"loc": dist.constraints.real_vector}  # noqa: RUF012
    support = dist.constraints.real_vector
    reparametrized_params = ["loc"]  # noqa: RUF012
    pytree_data_fields = (
        "loc",
        "precision_scale",
        "structure",
        "null_space",
        "precision_factors",
    )
    pytree_aux_fields = ("constraint", "soft_constraint_scale", "include_normalizer")

    def __init__(
        self,
        loc: Float[ArrayLike, " N"],
        precision_scale: Float[ArrayLike, ""],
        structure: lx.AbstractLinearOperator,
        null_space: Float[ArrayLike, "N c"] | Float[ArrayLike, " N"],
        precision_factors: tuple[lx.AbstractLinearOperator, ...] | None = None,
        constraint: Literal["none", "soft", "hard"] = "hard",
        soft_constraint_scale: float = 1e-3,
        include_normalizer: bool = False,
        *,
        validate_args: bool | None = None,
    ) -> None:
        if constraint not in ("none", "soft", "hard"):
            raise ValueError(
                f"constraint must be 'none', 'soft' or 'hard', got {constraint!r}."
            )
        loc = jnp.asarray(loc)
        n = structure.in_size()
        if loc.shape != (n,):
            raise ValueError(
                f"loc must have shape ({n},) to match the structure, got {loc.shape}."
            )
        null_space = jnp.asarray(null_space, dtype=loc.dtype)
        if null_space.ndim == 1:
            null_space = rearrange(null_space, "n -> n 1")
        if null_space.shape[0] != n:
            raise ValueError(
                f"null_space must have {n} rows, got shape {null_space.shape}."
            )
        self.loc = loc
        self.precision_scale = jnp.asarray(precision_scale, dtype=loc.dtype)
        self.structure = structure
        self.null_space = null_space
        self.precision_factors = (
            None if precision_factors is None else tuple(precision_factors)
        )
        self.constraint = constraint
        self.soft_constraint_scale = soft_constraint_scale
        self.include_normalizer = include_normalizer
        super().__init__(batch_shape=(), event_shape=(n,), validate_args=validate_args)

    # -- null space -----------------------------------------------------------

    def _null_dual(self) -> Float[Array, "N c"]:
        """``V (VᵀV)⁻¹``, so that ``P⊥x = x − V (VᵀV)⁻¹ Vᵀx``."""
        V = self.null_space
        gram = einsum(V, V, "n c, n d -> c d")
        return rearrange(
            jnp.linalg.solve(gram, rearrange(V, "n c -> c n")), "c n -> n c"
        )

    def _project(self, x: Float[Array, "*sample N"]) -> Float[Array, "*sample N"]:
        """``x − V (VᵀV)⁻¹ Vᵀx``, orthogonal to ``null_space``."""
        coeffs = einsum(x, self.null_space, "... n, n c -> ... c")
        return x - einsum(coeffs, self._null_dual(), "... c, n c -> ... n")

    def _shifted(self) -> lx.AbstractLinearOperator:
        """``τR + εI`` with ``ε = √eps · mean(diag(τR))``, as R-INLA does."""
        scaled = _scaled(self.structure, self.precision_scale)
        diagonal = diag(scaled)
        eps = float(jnp.finfo(diagonal.dtype).eps) ** 0.5 * jnp.mean(diagonal)
        return _add_ridge(scaled, eps)

    def _eigen(self) -> FactoredEigen | None:
        """``τR``'s factored eigendecomposition, for a Kronecker structure."""
        operator = _unwrap(self.structure)
        if not isinstance(operator, _EIGEN):
            return None
        eigen = factored_eigen(operator)
        if eigen is None:
            return None
        return FactoredEigen(eigen.bases, self.precision_scale * eigen.eigenvalues)

    def _sampler(self) -> tuple[int, _Draw]:
        """``(m, f)`` with ``P⊥ f(z) ~ N(0, (τR)⁺)`` for ``z ~ N(0, I_m)``."""
        n = self.event_shape[0]
        eigen = self._eigen()
        if eigen is not None:
            return n, lambda z: _eigen_draw(eigen, z, pinv=True)
        factors = self.precision_factors
        if factors is not None and not _is_factorable(_unwrap(self.structure)):
            root = jnp.sqrt(self.precision_scale)
            scaled = tuple(root * factor for factor in factors)
            return _perturbation_draw(
                _scaled(self.structure, self.precision_scale), scaled, None
            )
        return _precision_sampler(self._shifted(), None, None)

    # -- density --------------------------------------------------------------

    def _rank(self) -> int:
        return self.event_shape[0] - self.null_space.shape[1]

    def _half_log_pdet(self) -> Float[Array, ""]:
        """``½ log|R|₊``: eigenvalues of a Kronecker sum, else via the null space."""
        operator = _unwrap(self.structure)
        if isinstance(
            operator, KroneckerSum | DiagonalizedOperator | lx.DiagonalLinearOperator
        ):
            return 0.5 * pseudo_logdet(operator)
        return 0.5 * pseudo_logdet(operator, null_space=self.null_space)

    @validate_sample
    def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
        r"""$\tfrac{N-c}{2}\log\tau - \tfrac\tau2(x-\mu)^\top R(x-\mu)$, and extras.

        Args:
            value: Points, shape ``(..., N)``.

        Returns:
            Log-densities, shape ``(...)``.
        """
        tau = self.precision_scale
        rank = self._rank()
        quad = _quadratic(self.structure, value - self.loc)
        result = 0.5 * rank * jnp.log(tau) - 0.5 * tau * quad
        if self.include_normalizer:
            result = result + self._half_log_pdet() - 0.5 * rank * _LOG_2PI
        if self.constraint == "soft":
            scale = self.soft_constraint_scale
            coeffs = einsum(value, self.null_space, "... n, n c -> ... c")
            result = result - 0.5 * reduce(coeffs**2, "... c -> ...", "sum") / scale**2
            if self.include_normalizer:
                c = self.null_space.shape[1]
                result = result - c * (math.log(scale) + 0.5 * _LOG_2PI)
        return result

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*sample N"]:
        """Draws orthogonal to ``null_space`` (plus the soft or ``"none"`` part).

        Args:
            key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
            sample_shape: Leading sample shape.

        Returns:
            Draws, shape ``sample_shape + (N,)``. With ``"hard"``,
            ``null_spaceᵀ x = 0`` to machine precision.
        """
        if key is None:
            key = jax.random.PRNGKey(0)
        key_field, key_null = jax.random.split(key)
        size, draw = self._sampler()
        dtype = self.loc.dtype
        u = _draw_samples(key_field, sample_shape, size, draw, dtype)
        # Twice is enough: a shifted-Cholesky draw carries a ~ε^{-1/2} null
        # component, and one projection leaves eps times that.
        u = self._project(self._project(u))
        if self.constraint == "none":
            return self.loc + u
        x = self._project(self.loc) + u
        if self.constraint == "soft":
            c = self.null_space.shape[1]
            z = jax.random.normal(key_null, (*sample_shape, c), dtype=dtype)
            z = self.soft_constraint_scale * z
            x = x + einsum(z, self._null_dual(), "... c, n c -> ... n")
        return x

    def marginal_variances(self) -> Float[Array, " N"]:
        r"""``diag((τR)⁺)``, plus ``s² diag(V (VᵀV)⁻² Vᵀ)`` for ``"soft"``.

        The pseudo-inverse diagonal is exact from the factor eigenvalues on a
        Kronecker structure, and otherwise the kriging-corrected
        ``diag((τR + εI)⁻¹)`` (R-INLA's constrained marginal variances).

        Returns:
            Marginal variances of the samples, shape ``(N,)``.
        """
        eigen = self._eigen()
        if eigen is not None:
            variances = eigen.diag_inv(pinv=True)
        else:
            shifted = GaussianMRF(jnp.zeros_like(self.loc), self._shifted())
            kriged = shifted.condition_on_constraints(
                rearrange(self.null_space, "n c -> c n"),
                jnp.zeros(self.null_space.shape[1], dtype=self.loc.dtype),
            )
            variances = kriged.marginal_variances()
        if self.constraint == "soft":
            dual = self._null_dual()
            scale = self.soft_constraint_scale
            variances = variances + scale**2 * reduce(dual**2, "n c -> n", "sum")
        return variances

    @lazy_property
    def mean(self) -> Float[Array, " N"]:
        if self.constraint == "none":
            return self.loc
        return self._project(self.loc)

    @lazy_property
    def variance(self) -> Float[Array, " N"]:
        return self.marginal_variances()

log_prob(value: Float[Array, '*sample N']) -> Float[Array, '*sample']

\(\tfrac{N-c}{2}\log\tau - \tfrac\tau2(x-\mu)^\top R(x-\mu)\), and extras.

Parameters:

Name	Type	Description	Default
value	Float[Array, '*sample N']	Points, shape (..., N).	required

Returns:

Type	Description
Float[Array, '*sample']	Log-densities, shape (...).

Source code in src/gaussx/_distributions/_gmrf.py

@validate_sample
def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
    r"""$\tfrac{N-c}{2}\log\tau - \tfrac\tau2(x-\mu)^\top R(x-\mu)$, and extras.

    Args:
        value: Points, shape ``(..., N)``.

    Returns:
        Log-densities, shape ``(...)``.
    """
    tau = self.precision_scale
    rank = self._rank()
    quad = _quadratic(self.structure, value - self.loc)
    result = 0.5 * rank * jnp.log(tau) - 0.5 * tau * quad
    if self.include_normalizer:
        result = result + self._half_log_pdet() - 0.5 * rank * _LOG_2PI
    if self.constraint == "soft":
        scale = self.soft_constraint_scale
        coeffs = einsum(value, self.null_space, "... n, n c -> ... c")
        result = result - 0.5 * reduce(coeffs**2, "... c -> ...", "sum") / scale**2
        if self.include_normalizer:
            c = self.null_space.shape[1]
            result = result - c * (math.log(scale) + 0.5 * _LOG_2PI)
    return result

sample(key: jax.Array | None, sample_shape: tuple[int, ...] = ()) -> Float[Array, '*sample N']

Draws orthogonal to null_space (plus the soft or "none" part).

Parameters:

Name	Type	Description	Default
key	Array | None	PRNG key. None means jax.random.PRNGKey(0).	required
sample_shape	tuple[int, ...]	Leading sample shape.	()

Returns:

Type	Description
Float[Array, '*sample N']	Draws, shape sample_shape + (N,). With "hard",
Float[Array, '*sample N']	null_spaceᵀ x = 0 to machine precision.

Source code in src/gaussx/_distributions/_gmrf.py

def sample(
    self,
    key: jax.Array | None,
    sample_shape: tuple[int, ...] = (),
) -> Float[Array, "*sample N"]:
    """Draws orthogonal to ``null_space`` (plus the soft or ``"none"`` part).

    Args:
        key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
        sample_shape: Leading sample shape.

    Returns:
        Draws, shape ``sample_shape + (N,)``. With ``"hard"``,
        ``null_spaceᵀ x = 0`` to machine precision.
    """
    if key is None:
        key = jax.random.PRNGKey(0)
    key_field, key_null = jax.random.split(key)
    size, draw = self._sampler()
    dtype = self.loc.dtype
    u = _draw_samples(key_field, sample_shape, size, draw, dtype)
    # Twice is enough: a shifted-Cholesky draw carries a ~ε^{-1/2} null
    # component, and one projection leaves eps times that.
    u = self._project(self._project(u))
    if self.constraint == "none":
        return self.loc + u
    x = self._project(self.loc) + u
    if self.constraint == "soft":
        c = self.null_space.shape[1]
        z = jax.random.normal(key_null, (*sample_shape, c), dtype=dtype)
        z = self.soft_constraint_scale * z
        x = x + einsum(z, self._null_dual(), "... c, n c -> ... n")
    return x

marginal_variances() -> Float[Array, ' N']

diag((τR)⁺), plus s² diag(V (VᵀV)⁻² Vᵀ) for "soft".

The pseudo-inverse diagonal is exact from the factor eigenvalues on a Kronecker structure, and otherwise the kriging-corrected diag((τR + εI)⁻¹) (R-INLA's constrained marginal variances).

Returns:

Type	Description
Float[Array, ' N']	Marginal variances of the samples, shape (N,).

Source code in src/gaussx/_distributions/_gmrf.py

def marginal_variances(self) -> Float[Array, " N"]:
    r"""``diag((τR)⁺)``, plus ``s² diag(V (VᵀV)⁻² Vᵀ)`` for ``"soft"``.

    The pseudo-inverse diagonal is exact from the factor eigenvalues on a
    Kronecker structure, and otherwise the kriging-corrected
    ``diag((τR + εI)⁻¹)`` (R-INLA's constrained marginal variances).

    Returns:
        Marginal variances of the samples, shape ``(N,)``.
    """
    eigen = self._eigen()
    if eigen is not None:
        variances = eigen.diag_inv(pinv=True)
    else:
        shifted = GaussianMRF(jnp.zeros_like(self.loc), self._shifted())
        kriged = shifted.condition_on_constraints(
            rearrange(self.null_space, "n c -> c n"),
            jnp.zeros(self.null_space.shape[1], dtype=self.loc.dtype),
        )
        variances = kriged.marginal_variances()
    if self.constraint == "soft":
        dual = self._null_dual()
        scale = self.soft_constraint_scale
        variances = variances + scale**2 * reduce(dual**2, "n c -> n", "sum")
    return variances

ConstrainedGMRF

Bases: Distribution

A GaussianMRF conditioned on hard linear constraints \(A_cx = e\).

Conditioning by kriging (Rue & Held, 2005, §2.3.3): with \(W = Q^{-1}A_c^\top\) (c solves with one factorisation) and \(S = A_cW\), an unconstrained draw \(x\) maps to the exact constrained draw \(x^\ast = x - WS^{-1}(A_cx - e)\), so \(A_cx^\ast = e\) to machine precision. The marginal variances are \(\operatorname{diag}(Q^{-1}) - \operatorname{diag}(WS^{-1}W^\top)\), and the density (Rue & Held, eq. 2.30) is

\[ \log\pi(x \mid A_cx = e) = \log\pi(x) - \log\mathcal N(e;\,A_c\mu,\,S) - \tfrac12\log|A_cA_c^\top|, \]

a density on the affine subspace \(\{A_cx = e\}\) (w.r.t. its Lebesgue measure), evaluated at points assumed to satisfy the constraint.

Built by GaussianMRF.condition_on_constraints. Requires numpyro.

Parameters:

Name	Type	Description	Default
base	GaussianMRF	The unconstrained field.	required
constraint_matrix	Float[ArrayLike, 'c N'] | Float[ArrayLike, ' N']	\(A_c\), shape (c, N) (or (N,)), full row rank.	required
constraint_value	Float[ArrayLike, ' c'] | Float[ArrayLike, '']	\(e\), shape (c,).	required
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

import jax
import jax.numpy as jnp
import lineax as lx
import gaussx as gx

Q = lx.MatrixLinearOperator(2.0 * jnp.eye(4), lx.positive_semidefinite_tag)
prior = gx.GaussianMRF(jnp.zeros(4), Q)
field = prior.condition_on_constraints(jnp.ones(4), 1.0)  # Σ x = 1
x = field.sample(jax.random.PRNGKey(0), (10,))  # each row sums to 1
var = field.marginal_variances()  # 0.5 - 0.5 / 4

Source code in src/gaussx/_distributions/_gmrf.py

class ConstrainedGMRF(dist.Distribution):
    r"""A `GaussianMRF` conditioned on hard linear constraints $A_cx = e$.

    Conditioning by kriging (Rue & Held, 2005, §2.3.3): with
    $W = Q^{-1}A_c^\top$ (``c`` solves with one factorisation) and
    $S = A_cW$, an unconstrained draw $x$ maps to the exact constrained draw
    $x^\ast = x - WS^{-1}(A_cx - e)$, so $A_cx^\ast = e$ to machine
    precision. The marginal variances are
    $\operatorname{diag}(Q^{-1}) - \operatorname{diag}(WS^{-1}W^\top)$, and
    the density (Rue & Held, eq. 2.30) is

    $$
    \log\pi(x \mid A_cx = e) = \log\pi(x) - \log\mathcal N(e;\,A_c\mu,\,S)
        - \tfrac12\log|A_cA_c^\top|,
    $$

    a density on the affine subspace $\{A_cx = e\}$ (w.r.t. its Lebesgue
    measure), evaluated at points assumed to satisfy the constraint.

    Built by `GaussianMRF.condition_on_constraints`. Requires ``numpyro``.

    Args:
        base: The unconstrained field.
        constraint_matrix: $A_c$, shape ``(c, N)`` (or ``(N,)``), full row rank.
        constraint_value: $e$, shape ``(c,)``.
        validate_args: Whether to validate input arguments.

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

        Q = lx.MatrixLinearOperator(2.0 * jnp.eye(4), lx.positive_semidefinite_tag)
        prior = gx.GaussianMRF(jnp.zeros(4), Q)
        field = prior.condition_on_constraints(jnp.ones(4), 1.0)  # Σ x = 1
        x = field.sample(jax.random.PRNGKey(0), (10,))  # each row sums to 1
        var = field.marginal_variances()  # 0.5 - 0.5 / 4
        ```
    """

    arg_constraints = {}  # noqa: RUF012
    support = dist.constraints.real_vector
    pytree_data_fields = ("base", "constraint_matrix", "constraint_value")

    def __init__(
        self,
        base: GaussianMRF,
        constraint_matrix: Float[ArrayLike, "c N"] | Float[ArrayLike, " N"],
        constraint_value: Float[ArrayLike, " c"] | Float[ArrayLike, ""],
        *,
        validate_args: bool | None = None,
    ) -> None:
        dtype = base.loc.dtype
        A = jnp.atleast_2d(jnp.asarray(constraint_matrix, dtype=dtype))
        e = jnp.atleast_1d(jnp.asarray(constraint_value, dtype=dtype))
        n = base.event_shape[0]
        if A.shape[1] != n or e.shape != (A.shape[0],):
            raise ValueError(
                f"constraint_matrix must have shape (c, {n}) and constraint_value "
                f"shape (c,), got {A.shape} and {e.shape}."
            )
        self.base = base
        self.constraint_matrix = A
        self.constraint_value = e
        super().__init__(batch_shape=(), event_shape=(n,), validate_args=validate_args)

    def _kriging(self) -> tuple[Float[Array, "c N"], Float[Array, "c N"], Array]:
        """``W = Q⁻¹A_cᵀ`` (rows), ``K = S⁻¹Wᵀ`` (rows) and ``chol(S)``."""
        A = self.constraint_matrix
        W = jax.vmap(_solve_fn(self.base.precision, self.base.solver))(A)
        S = einsum(A, W, "c n, d n -> c d")
        chol_S = jnp.linalg.cholesky(S)
        K = jsl.cho_solve((chol_S, True), W)
        return W, K, chol_S

    def _correct(self, x: Float[Array, "*sample N"], K: Array) -> Array:
        """``x − W S⁻¹ (A_c x − e)``."""
        flat = _flat(x)
        excess = einsum(flat, self.constraint_matrix, "s n, c n -> s c")
        excess = excess - self.constraint_value
        corrected = flat - einsum(excess, K, "s c, c n -> s n")
        return _reshape_samples(corrected, x.shape[:-1])

    @validate_sample
    def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
        """The constrained density (Rue & Held, eq. 2.30) on ``A_c x = e``.

        Args:
            value: Points satisfying the constraint, shape ``(..., N)``.

        Returns:
            Log-densities, shape ``(...)``.
        """
        A, e = self.constraint_matrix, self.constraint_value
        _, _, chol_S = self._kriging()
        c = A.shape[0]
        residual = e - einsum(A, self.base.loc, "c n, n -> c")
        white = jsl.solve_triangular(chol_S, residual, lower=True)
        log_det_S = 2.0 * jnp.sum(jnp.log(jnp.diagonal(chol_S)))
        log_constraint = -0.5 * (
            c * _LOG_2PI + log_det_S + einsum(white, white, "c, c ->")
        )
        gram = einsum(A, A, "c n, d n -> c d")
        log_det_gram = jnp.linalg.slogdet(gram)[1]
        return self.base.log_prob(value) - log_constraint - 0.5 * log_det_gram

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*sample N"]:
        """Base draws, kriged onto ``A_c x = e``.

        Args:
            key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
            sample_shape: Leading sample shape.

        Returns:
            Draws, shape ``sample_shape + (N,)``.
        """
        _, K, _ = self._kriging()
        return self._correct(self.base.sample(key, sample_shape), K)

    def marginal_variances(self) -> Float[Array, " N"]:
        """``diag(Q⁻¹) − diag(W S⁻¹ Wᵀ)``.

        Returns:
            Constrained marginal variances, shape ``(N,)``.
        """
        W, K, _ = self._kriging()
        correction = reduce(W * K, "c n -> n", "sum")
        return self.base.marginal_variances() - correction

    @lazy_property
    def mean(self) -> Float[Array, " N"]:
        _, K, _ = self._kriging()
        return self._correct(self.base.loc, K)

    @lazy_property
    def variance(self) -> Float[Array, " N"]:
        return self.marginal_variances()

log_prob(value: Float[Array, '*sample N']) -> Float[Array, '*sample']

The constrained density (Rue & Held, eq. 2.30) on A_c x = e.

Parameters:

Name	Type	Description	Default
value	Float[Array, '*sample N']	Points satisfying the constraint, shape (..., N).	required

Returns:

Type	Description
Float[Array, '*sample']	Log-densities, shape (...).

Source code in src/gaussx/_distributions/_gmrf.py

@validate_sample
def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
    """The constrained density (Rue & Held, eq. 2.30) on ``A_c x = e``.

    Args:
        value: Points satisfying the constraint, shape ``(..., N)``.

    Returns:
        Log-densities, shape ``(...)``.
    """
    A, e = self.constraint_matrix, self.constraint_value
    _, _, chol_S = self._kriging()
    c = A.shape[0]
    residual = e - einsum(A, self.base.loc, "c n, n -> c")
    white = jsl.solve_triangular(chol_S, residual, lower=True)
    log_det_S = 2.0 * jnp.sum(jnp.log(jnp.diagonal(chol_S)))
    log_constraint = -0.5 * (
        c * _LOG_2PI + log_det_S + einsum(white, white, "c, c ->")
    )
    gram = einsum(A, A, "c n, d n -> c d")
    log_det_gram = jnp.linalg.slogdet(gram)[1]
    return self.base.log_prob(value) - log_constraint - 0.5 * log_det_gram

sample(key: jax.Array | None, sample_shape: tuple[int, ...] = ()) -> Float[Array, '*sample N']

Base draws, kriged onto A_c x = e.

Parameters:

Name	Type	Description	Default
key	Array | None	PRNG key. None means jax.random.PRNGKey(0).	required
sample_shape	tuple[int, ...]	Leading sample shape.	()

Returns:

Type	Description
Float[Array, '*sample N']	Draws, shape sample_shape + (N,).

Source code in src/gaussx/_distributions/_gmrf.py

def sample(
    self,
    key: jax.Array | None,
    sample_shape: tuple[int, ...] = (),
) -> Float[Array, "*sample N"]:
    """Base draws, kriged onto ``A_c x = e``.

    Args:
        key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
        sample_shape: Leading sample shape.

    Returns:
        Draws, shape ``sample_shape + (N,)``.
    """
    _, K, _ = self._kriging()
    return self._correct(self.base.sample(key, sample_shape), K)

marginal_variances() -> Float[Array, ' N']

diag(Q⁻¹) − diag(W S⁻¹ Wᵀ).

Returns:

Type	Description
Float[Array, ' N']	Constrained marginal variances, shape (N,).

Source code in src/gaussx/_distributions/_gmrf.py

def marginal_variances(self) -> Float[Array, " N"]:
    """``diag(Q⁻¹) − diag(W S⁻¹ Wᵀ)``.

    Returns:
        Constrained marginal variances, shape ``(N,)``.
    """
    W, K, _ = self._kriging()
    correction = reduce(W * K, "c n -> n", "sum")
    return self.base.marginal_variances() - correction

BYM2GMRF

Bases: IntrinsicGMRF

The BYM2 pair \((b, u^\ast)\) (Riebler et al., 2016), with its exact density.

\(b = (\sqrt{1-\phi}\,v + \sqrt\phi\,u^\ast)/\sqrt\tau\) with \(v \sim \mathcal N(0, I)\) and \(u^\ast\) the scaled ICAR field (structure \(R^\ast\), generalized_variance_scale) under one sum-to-zero constraint per connected component. The joint precision \(Q(\tau, \phi)\) is bym2_precision, and the field lives on \(\{V^\top u^\ast = 0\}\), where its log-density (Rue & Held, eq. 2.30) is

\[ \log\pi(b, u^\ast) = -\tfrac12 x^\top Q x + \tfrac n2 \log\tfrac{\tau}{1-\phi} \;\Big[+\ \tfrac12\log|R^\ast|_+ - \tfrac{2n-c}{2}\log 2\pi\Big], \]

because the Schur complement of \(Q\)'s \(b\)-block is exactly \(R^\ast\).

Why not IntrinsicGMRF on \(Q\): \(Q\)'s null vector is \((\sqrt{\phi/\tau}\,\mathbf 1, \mathbf 1)\), which depends on \(\theta\). So the bracket IntrinsicGMRF drops by default is not \(\theta\)-free (it holds the \(\tfrac n2\log\tfrac{\tau}{1-\phi}\) above), its normaliser takes a dense determinant of size \(2n\) at every \(\theta\), and its sampler's projection along null_space leaves the \(b\)-part of a shifted-Cholesky draw unbounded. Here the \(\theta\)-dependent term is always included, the bracket is truly constant (from \(R^\ast\) alone: the matrix-tree theorem, one sparse Cholesky, or log_pdet), and draws are exact: \(u^\ast\) from the constrained ICAR on \(R^\ast\), then \(b\) from its definition.

It is an IntrinsicGMRF whose structure is \(Q\), precision_scale is 1 and null_space is \(V\) padded with zeros on the \(b\) rows, so laplace_mode and vb_mean_correction take it as a prior unchanged (their log-marginal then has the correct \(\theta\)-dependence).

Parameters:

Name	Type	Description	Default
structure_scaled	AbstractLinearOperator	\(R^\ast = sR\), a (scaled) SparseOperator graph Laplacian of size n.	required
tau	Float[ArrayLike, '']	Precision \(\tau > 0\) of \(b\) (may be traced).	required
phi	Float[ArrayLike, '']	Mixing \(0 \le \phi < 1\) (may be traced).	required
null_space	Float[ArrayLike, 'n c'] | Float[ArrayLike, ' n']	Basis of \(\ker R^\ast\), shape (n, c) (or (n,)): one column per connected component (kernellib's graph_null_space).	required
loc	Float[ArrayLike, ' N'] | None	Mean of the stacked (b, u*), shape (2n,); zeros by default.	None
constraint	Literal['none', 'soft', 'hard']	"hard" (default), "soft" or "none", on \(u^\ast\) as in IntrinsicGMRF.	'hard'
soft_constraint_scale	float	\(s\) for "soft".	0.001
include_normalizer	bool	Add the \(\theta\)-free constant in brackets.	False
log_pdet	Float[ArrayLike, ''] | None	\(\log|R^\ast|_+\), if already known: it is \(\theta\)-free, so compute it once outside a \(\theta\) loop. None computes it per call by the matrix-tree theorem.	None
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

import jax
import jax.numpy as jnp
import numpy as np
import gaussx as gx

# A path of 3 areas: R = D - W
R = gx.SparseOperator.from_coo(
    np.array([0, 1, 2, 1, 2]),
    np.array([0, 1, 2, 0, 1]),
    jnp.array([1.0, 2.0, 1.0, -1.0, -1.0]),
    (3, 3),
    symmetric=True,
)
V = jnp.ones(3) / jnp.sqrt(3.0)
R_star = gx.generalized_variance_scale(R, V) * R
prior = gx.BYM2GMRF(R_star, tau=2.0, phi=0.6, null_space=V)
x = prior.sample(jax.random.PRNGKey(0), (100,))  # (100, 6): (b, u*)
lp = prior.log_prob(x)  # correct in (tau, phi) for a marginal likelihood

Source code in src/gaussx/_distributions/_gmrf.py

class BYM2GMRF(IntrinsicGMRF):
    r"""The BYM2 pair $(b, u^\ast)$ (Riebler et al., 2016), with its exact density.

    $b = (\sqrt{1-\phi}\,v + \sqrt\phi\,u^\ast)/\sqrt\tau$ with
    $v \sim \mathcal N(0, I)$ and $u^\ast$ the scaled ICAR field (structure
    $R^\ast$, `generalized_variance_scale`) under one sum-to-zero constraint
    per connected component. The joint precision $Q(\tau, \phi)$ is
    `bym2_precision`, and the field lives on $\{V^\top u^\ast = 0\}$, where
    its log-density (Rue & Held, eq. 2.30) is

    $$
    \log\pi(b, u^\ast) = -\tfrac12 x^\top Q x
        + \tfrac n2 \log\tfrac{\tau}{1-\phi}
        \;\Big[+\ \tfrac12\log|R^\ast|_+ - \tfrac{2n-c}{2}\log 2\pi\Big],
    $$

    because the Schur complement of $Q$'s $b$-block is exactly $R^\ast$.

    Why not `IntrinsicGMRF` on $Q$: $Q$'s null vector is
    $(\sqrt{\phi/\tau}\,\mathbf 1, \mathbf 1)$, which depends on $\theta$.
    So the bracket `IntrinsicGMRF` drops by default is not $\theta$-free (it
    holds the $\tfrac n2\log\tfrac{\tau}{1-\phi}$ above), its normaliser
    takes a dense determinant of size $2n$ at every $\theta$, and its
    sampler's projection along ``null_space`` leaves the $b$-part of a
    shifted-Cholesky draw unbounded. Here the $\theta$-dependent term is
    always included, the bracket is truly constant (from $R^\ast$ alone:
    the matrix-tree theorem, one sparse Cholesky, or ``log_pdet``), and
    draws are exact: $u^\ast$ from the constrained ICAR on $R^\ast$, then
    $b$ from its definition.

    It is an `IntrinsicGMRF` whose ``structure`` is $Q$, ``precision_scale``
    is 1 and ``null_space`` is $V$ padded with zeros on the $b$ rows, so
    `laplace_mode` and `vb_mean_correction` take it as a prior unchanged
    (their log-marginal then has the correct $\theta$-dependence).

    Args:
        structure_scaled: $R^\ast = sR$, a (scaled) `SparseOperator` graph
            Laplacian of size ``n``.
        tau: Precision $\tau > 0$ of $b$ (may be traced).
        phi: Mixing $0 \le \phi < 1$ (may be traced).
        null_space: Basis of $\ker R^\ast$, shape ``(n, c)`` (or ``(n,)``):
            one column per connected component (kernellib's
            ``graph_null_space``).
        loc: Mean of the stacked ``(b, u*)``, shape ``(2n,)``; zeros by
            default.
        constraint: ``"hard"`` (default), ``"soft"`` or ``"none"``, on
            $u^\ast$ as in `IntrinsicGMRF`.
        soft_constraint_scale: $s$ for ``"soft"``.
        include_normalizer: Add the $\theta$-free constant in brackets.
        log_pdet: $\log|R^\ast|_+$, if already known: it is
            $\theta$-free, so compute it once outside a $\theta$ loop.
            ``None`` computes it per call by the matrix-tree theorem.
        validate_args: Whether to validate input arguments.

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

        # A path of 3 areas: R = D - W
        R = gx.SparseOperator.from_coo(
            np.array([0, 1, 2, 1, 2]),
            np.array([0, 1, 2, 0, 1]),
            jnp.array([1.0, 2.0, 1.0, -1.0, -1.0]),
            (3, 3),
            symmetric=True,
        )
        V = jnp.ones(3) / jnp.sqrt(3.0)
        R_star = gx.generalized_variance_scale(R, V) * R
        prior = gx.BYM2GMRF(R_star, tau=2.0, phi=0.6, null_space=V)
        x = prior.sample(jax.random.PRNGKey(0), (100,))  # (100, 6): (b, u*)
        lp = prior.log_prob(x)  # correct in (tau, phi) for a marginal likelihood
        ```
    """

    pytree_data_fields = (
        "tau",
        "phi",
        "structure_scaled",
        "icar_null_space",
        "log_pdet",
    )
    pytree_aux_fields = ()

    def __init__(
        self,
        structure_scaled: lx.AbstractLinearOperator,
        tau: Float[ArrayLike, ""],
        phi: Float[ArrayLike, ""],
        null_space: Float[ArrayLike, "n c"] | Float[ArrayLike, " n"],
        loc: Float[ArrayLike, " N"] | None = None,
        *,
        constraint: Literal["none", "soft", "hard"] = "hard",
        soft_constraint_scale: float = 1e-3,
        include_normalizer: bool = False,
        log_pdet: Float[ArrayLike, ""] | None = None,
        validate_args: bool | None = None,
    ) -> None:
        n = structure_scaled.in_size()
        precision = bym2_precision(structure_scaled, tau, phi)
        dtype = precision.values.dtype
        if loc is None:
            loc = jnp.zeros(2 * n, dtype=dtype)
        V = jnp.asarray(null_space, dtype=dtype)
        if V.ndim == 1:
            V = rearrange(V, "n -> n 1")
        if V.shape[0] != n:
            raise ValueError(f"null_space must have {n} rows, got shape {V.shape}.")
        padded = jnp.concatenate([jnp.zeros_like(V), V])
        super().__init__(
            loc,
            1.0,
            precision,
            padded,
            constraint=constraint,
            soft_constraint_scale=soft_constraint_scale,
            include_normalizer=include_normalizer,
            validate_args=validate_args,
        )
        self.tau = jnp.asarray(tau, dtype=dtype)
        self.phi = jnp.asarray(phi, dtype=dtype)
        self.structure_scaled = structure_scaled
        self.icar_null_space = V
        self.log_pdet = None if log_pdet is None else jnp.asarray(log_pdet, dtype)

    def _n(self) -> int:
        return self.icar_null_space.shape[0]

    def _icar(self, constraint: str) -> IntrinsicGMRF:
        """The unit-precision ICAR on ``R*`` that ``u*`` is drawn from."""
        n = self._n()
        return IntrinsicGMRF(
            jnp.zeros(n, dtype=self.loc.dtype),
            1.0,
            self.structure_scaled,
            self.icar_null_space,
            constraint=constraint,
            soft_constraint_scale=self.soft_constraint_scale,
        )

    def _half_log_pdet(self) -> Float[Array, ""]:
        """``½ log|R*|₊``: given, or by the matrix-tree theorem (one sparse
        Cholesky)."""
        if self.log_pdet is not None:
            return 0.5 * self.log_pdet
        values, pattern = _as_sparse(self.structure_scaled)
        laplacian = SparseOperator(values, pattern)
        return 0.5 * pseudo_logdet(laplacian, structure="laplacian")

    @validate_sample
    def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
        r"""$-\tfrac12 x^\top Q x + \tfrac n2\log\tfrac{\tau}{1-\phi}$, and extras.

        Args:
            value: Points ``(b, u*)``, shape ``(..., 2n)``.

        Returns:
            Log-densities, shape ``(...)``.
        """
        n = self._n()
        quad = _quadratic(self.structure, value - self.loc)
        result = -0.5 * quad + 0.5 * n * jnp.log(self.tau / (1.0 - self.phi))
        if self.include_normalizer:
            rank = 2 * n - self.icar_null_space.shape[1]
            result = result + self._half_log_pdet() - 0.5 * rank * _LOG_2PI
        if self.constraint == "soft":
            scale = self.soft_constraint_scale
            coeffs = einsum(value, self.null_space, "... n, n c -> ... c")
            result = result - 0.5 * reduce(coeffs**2, "... c -> ...", "sum") / scale**2
            if self.include_normalizer:
                c = self.null_space.shape[1]
                result = result - c * (math.log(scale) + 0.5 * _LOG_2PI)
        return result

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*sample N"]:
        """Exact draws: ``u*`` from the ICAR on ``R*``, then ``b`` from ``u*``.

        Args:
            key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
            sample_shape: Leading sample shape.

        Returns:
            Draws ``(b, u*)``, shape ``sample_shape + (2n,)``.
        """
        if key is None:
            key = jax.random.PRNGKey(0)
        key_u, key_v = jax.random.split(key)
        n = self._n()
        u = self._icar(self.constraint).sample(key_u, sample_shape)
        v = jax.random.normal(key_v, (*sample_shape, n), dtype=self.loc.dtype)
        b = (jnp.sqrt(1.0 - self.phi) * v + jnp.sqrt(self.phi) * u) / jnp.sqrt(self.tau)
        centre = self.loc if self.constraint == "none" else self._project(self.loc)
        return centre + jnp.concatenate([b, u], axis=-1)

    def marginal_variances(self) -> Float[Array, " N"]:
        r"""``var(u*)`` from the ICAR on ``R*``, ``var(b) = ((1−φ) + φ var(u*))/τ``.

        Returns:
            Marginal variances of ``(b, u*)``, shape ``(2n,)``.
        """
        var_u = self._icar(self.constraint).marginal_variances()
        var_b = ((1.0 - self.phi) + self.phi * var_u) / self.tau
        return jnp.concatenate([var_b, var_u])

log_prob(value: Float[Array, '*sample N']) -> Float[Array, '*sample']

\(-\tfrac12 x^\top Q x + \tfrac n2\log\tfrac{\tau}{1-\phi}\), and extras.

Parameters:

Name	Type	Description	Default
value	Float[Array, '*sample N']	Points (b, u*), shape (..., 2n).	required

Returns:

Type	Description
Float[Array, '*sample']	Log-densities, shape (...).

Source code in src/gaussx/_distributions/_gmrf.py

@validate_sample
def log_prob(self, value: Float[Array, "*sample N"]) -> Float[Array, "*sample"]:
    r"""$-\tfrac12 x^\top Q x + \tfrac n2\log\tfrac{\tau}{1-\phi}$, and extras.

    Args:
        value: Points ``(b, u*)``, shape ``(..., 2n)``.

    Returns:
        Log-densities, shape ``(...)``.
    """
    n = self._n()
    quad = _quadratic(self.structure, value - self.loc)
    result = -0.5 * quad + 0.5 * n * jnp.log(self.tau / (1.0 - self.phi))
    if self.include_normalizer:
        rank = 2 * n - self.icar_null_space.shape[1]
        result = result + self._half_log_pdet() - 0.5 * rank * _LOG_2PI
    if self.constraint == "soft":
        scale = self.soft_constraint_scale
        coeffs = einsum(value, self.null_space, "... n, n c -> ... c")
        result = result - 0.5 * reduce(coeffs**2, "... c -> ...", "sum") / scale**2
        if self.include_normalizer:
            c = self.null_space.shape[1]
            result = result - c * (math.log(scale) + 0.5 * _LOG_2PI)
    return result

sample(key: jax.Array | None, sample_shape: tuple[int, ...] = ()) -> Float[Array, '*sample N']

Exact draws: u* from the ICAR on R*, then b from u*.

Parameters:

Name	Type	Description	Default
key	Array | None	PRNG key. None means jax.random.PRNGKey(0).	required
sample_shape	tuple[int, ...]	Leading sample shape.	()

Returns:

Type	Description
Float[Array, '*sample N']	Draws (b, u*), shape sample_shape + (2n,).

Source code in src/gaussx/_distributions/_gmrf.py

def sample(
    self,
    key: jax.Array | None,
    sample_shape: tuple[int, ...] = (),
) -> Float[Array, "*sample N"]:
    """Exact draws: ``u*`` from the ICAR on ``R*``, then ``b`` from ``u*``.

    Args:
        key: PRNG key. ``None`` means ``jax.random.PRNGKey(0)``.
        sample_shape: Leading sample shape.

    Returns:
        Draws ``(b, u*)``, shape ``sample_shape + (2n,)``.
    """
    if key is None:
        key = jax.random.PRNGKey(0)
    key_u, key_v = jax.random.split(key)
    n = self._n()
    u = self._icar(self.constraint).sample(key_u, sample_shape)
    v = jax.random.normal(key_v, (*sample_shape, n), dtype=self.loc.dtype)
    b = (jnp.sqrt(1.0 - self.phi) * v + jnp.sqrt(self.phi) * u) / jnp.sqrt(self.tau)
    centre = self.loc if self.constraint == "none" else self._project(self.loc)
    return centre + jnp.concatenate([b, u], axis=-1)

marginal_variances() -> Float[Array, ' N']

var(u*) from the ICAR on R*, var(b) = ((1−φ) + φ var(u*))/τ.

Returns:

Type	Description
Float[Array, ' N']	Marginal variances of (b, u*), shape (2n,).

Source code in src/gaussx/_distributions/_gmrf.py

def marginal_variances(self) -> Float[Array, " N"]:
    r"""``var(u*)`` from the ICAR on ``R*``, ``var(b) = ((1−φ) + φ var(u*))/τ``.

    Returns:
        Marginal variances of ``(b, u*)``, shape ``(2n,)``.
    """
    var_u = self._icar(self.constraint).marginal_variances()
    var_b = ((1.0 - self.phi) + self.phi * var_u) / self.tau
    return jnp.concatenate([var_b, var_u])

Sequential distributions

The linear-Gaussian state-space model as a density rather than a set of functions: log_prob is the Kalman marginal likelihood \(\log p(y_{1:T})\) (delegated to kalman_filter), sample is ancestral forward simulation, and event_shape is \((T, M)\) so log_prob returns a scalar with no .to_event() wrapping. That makes them usable directly as a NumPyro likelihood site.

A, H, Q and R each take a dense array or a lineax operator, exactly matching kalman_filter's contract — a structured \(Q\) / \(R\) keeps its structure through the Cholesky in sample, and \(H\) through the sandwich in variance.

MaskedLGSSM carries a \((T, M)\) observation mask and returns the exact marginal \(\log p(y_{\mathrm{obs}})\) — not a bound — because \(p(y_{\mathrm{miss}} \mid y_{\mathrm{obs}})\) is closed-form Gaussian. LGSSMFactory is the mask -> MaskedLGSSM callable for conditional use; it is an equinox.Module rather than a closure so the state-space parameters stay visible to equinox.filter_grad.

To use one as a normalizing-flow base, wrap it with gauss_flows.NumpyroBase, which already adapts any numpyro distribution — no bespoke adapter class is needed on either side. All three require numpyro to be installed.

Structured linear algebra and Gaussian primitives for JAX.

LGSSM

Bases: Distribution

Linear-Gaussian state-space model as a density over (T, M).

Models

\[ x_t = A x_{t-1} + q_t,\quad q_t \sim N(0, Q), \qquad y_t = H x_t + r_t,\quad r_t \sim N(0, R), \]

with \(x_0 \sim N(m_0, P_0)\) and \(t = 1, \dots, T\). log_prob is the Kalman marginal likelihood \(\log p(y_{1:T})\), delegated to gaussx.kalman_filter; sample is ancestral forward simulation. event_shape is (T, M), so log_prob returns a scalar and no .to_event() wrapping is needed.

Requires the numpyro optional extra (pip install "gaussx[numpyro]").

A, H, Q and R each accept a dense array or a lineax.AbstractLinearOperator, exactly matching gaussx.kalman_filter's contract. Structured Q / R keep their structure through the Cholesky in sample, H through the sandwich in variance, and A / H apply via structural matvec rather than a dense matmul. m0 and P0 are dense, as they are for gaussx.kalman_filter.

Parameters:

Name	Type	Description	Default
A	_Transition	Transition matrix or operator, shape (N, N).	required
H	_Emission	Emission matrix or operator, shape (M, N).	required
Q	_StateCov	Process noise covariance or operator, shape (N, N).	required
R	_ObsCov	Observation noise covariance or operator, shape (M, M).	required
m0	Float[Array, ' N']	Initial state mean, shape (N,).	required
P0	Float[Array, 'N N']	Initial state covariance, shape (N, N).	required
n_steps	int	Sequence length T.	required
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from gaussx import LGSSM
>>> d = LGSSM(0.9 * jnp.eye(2), jnp.eye(2), 0.1 * jnp.eye(2),
...           0.2 * jnp.eye(2), jnp.zeros(2), jnp.eye(2), n_steps=10)
>>> y = d.sample(jr.key(0))
>>> y.shape
(10, 2)
>>> d.log_prob(y).shape
()

Notes

To use as a normalizing-flow base, wrap with gauss_flows.NumpyroBase; no bespoke adapter is required.

Source code in src/gaussx/_distributions/_lgssm.py

class LGSSM(dist.Distribution):
    r"""Linear-Gaussian state-space model as a density over ``(T, M)``.

    Models

    $$
    x_t = A x_{t-1} + q_t,\quad q_t \sim N(0, Q), \qquad
    y_t = H x_t + r_t,\quad r_t \sim N(0, R),
    $$

    with $x_0 \sim N(m_0, P_0)$ and $t = 1, \dots, T$. ``log_prob`` is
    the Kalman marginal likelihood $\log p(y_{1:T})$, delegated to
    `gaussx.kalman_filter`; ``sample`` is ancestral forward
    simulation. ``event_shape`` is ``(T, M)``, so ``log_prob`` returns a
    scalar and no ``.to_event()`` wrapping is needed.

    Requires the ``numpyro`` optional extra
    (``pip install "gaussx[numpyro]"``).

    ``A``, ``H``, ``Q`` and ``R`` each accept a dense array *or* a
    `lineax.AbstractLinearOperator`, exactly matching
    `gaussx.kalman_filter`'s contract. Structured ``Q`` / ``R`` keep
    their structure through the Cholesky in ``sample``, ``H`` through
    the sandwich in ``variance``, and ``A`` / ``H`` apply via
    structural matvec rather than a dense matmul. ``m0`` and ``P0`` are
    dense, as they are for `gaussx.kalman_filter`.

    Args:
        A: Transition matrix or operator, shape ``(N, N)``.
        H: Emission matrix or operator, shape ``(M, N)``.
        Q: Process noise covariance or operator, shape ``(N, N)``.
        R: Observation noise covariance or operator, shape ``(M, M)``.
        m0: Initial state mean, shape ``(N,)``.
        P0: Initial state covariance, shape ``(N, N)``.
        n_steps: Sequence length ``T``.
        validate_args: Whether to validate input arguments.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from gaussx import LGSSM
        >>> d = LGSSM(0.9 * jnp.eye(2), jnp.eye(2), 0.1 * jnp.eye(2),
        ...           0.2 * jnp.eye(2), jnp.zeros(2), jnp.eye(2), n_steps=10)
        >>> y = d.sample(jr.key(0))
        >>> y.shape
        (10, 2)
        >>> d.log_prob(y).shape
        ()

    Notes:
        To use as a normalizing-flow base, wrap with
        ``gauss_flows.NumpyroBase``; no bespoke adapter is required.
    """

    arg_constraints = {}  # noqa: RUF012
    support = dist.constraints.real_matrix
    reparametrized_params = ["A", "H", "Q", "R", "m0", "P0"]  # noqa: RUF012
    pytree_data_fields = ("A", "H", "Q", "R", "m0", "P0")

    def __init__(
        self,
        A: _Transition,
        H: _Emission,
        Q: _StateCov,
        R: _ObsCov,
        m0: Float[Array, " N"],
        P0: Float[Array, "N N"],
        n_steps: int,
        *,
        validate_args: bool | None = None,
    ) -> None:
        self.A = A
        self.H = H
        self.Q = Q
        self.R = R
        self.m0 = m0
        self.P0 = P0
        super().__init__(
            batch_shape=(),
            event_shape=(n_steps, _out_size(H)),
            validate_args=validate_args,
        )

    @property
    def n_steps(self) -> int:
        """Sequence length ``T``."""
        return self.event_shape[0]

    @property
    def _obs_mask(self) -> Bool[Array, "T M"] | None:
        """Observation mask handed to `gaussx.kalman_filter`."""
        return None

    def _log_prob_single(self, value: Float[Array, "T M"]) -> Float[Array, ""]:
        return kalman_filter(
            self.A,
            self.H,
            self.Q,
            self.R,
            value,
            self.m0,
            self.P0,
            mask=self._obs_mask,
        ).log_likelihood

    @validate_sample
    def log_prob(self, value: Float[Array, "*batch T M"]) -> Float[Array, "*batch"]:
        leading_shape = value.shape[:-2]
        value_flat = rearrange(value, "... T M -> (...) T M")
        log_prob_flat = jax.vmap(self._log_prob_single)(value_flat)
        return _reshape_batch(log_prob_flat, leading_shape)

    def _sample_single(self, key: jax.Array) -> Float[Array, "T M"]:
        """One ancestral forward simulation of the full ``(T, M)`` series.

        Cholesky factors go through the structural
        `gaussx.cholesky` primitive, so a Kronecker / block-diagonal
        / diagonal covariance is never materialised.
        """
        N = self.m0.shape[0]
        M = self.event_shape[1]
        key_init, key_q, key_r = jax.random.split(key, 3)
        L_P0 = _cholesky(_psd(self.P0))
        L_Q = _cholesky(_psd(self.Q))
        L_R = _cholesky(_psd(self.R))

        # Each noise term is drawn in its factor's dtype, not JAX's default
        # float, so a float32 model stays float32 under x64.
        x0 = self.m0 + L_P0.mv(
            jax.random.normal(key_init, (N,), dtype=L_P0.in_structure().dtype)
        )
        eps_q = jax.random.normal(
            key_q, (self.n_steps, N), dtype=L_Q.in_structure().dtype
        )
        eps_r = jax.random.normal(
            key_r, (self.n_steps, M), dtype=L_R.in_structure().dtype
        )

        def step(x, noise):
            e_q, e_r = noise
            x_next = _matvec(self.A, x) + L_Q.mv(e_q)
            y = _matvec(self.H, x_next) + L_R.mv(e_r)
            return x_next, y

        _, ys = jax.lax.scan(step, x0, (eps_q, eps_r))
        return ys

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*sample T M"]:
        if key is None:
            raise ValueError("PRNG key must be provided to sample from LGSSM.")
        n_samples = math.prod(sample_shape) if sample_shape else 1
        keys = jax.random.split(key, n_samples)
        samples_flat = jax.vmap(self._sample_single)(keys)
        return _reshape_series(samples_flat, sample_shape)

    @lazy_property
    def mean(self) -> Float[Array, "T M"]:
        """Marginal mean $H A^t m_0$ of each observation, shape ``(T, M)``."""

        def step(m, _):
            m_next = _matvec(self.A, m)
            return m_next, _matvec(self.H, m_next)

        _, means = jax.lax.scan(step, self.m0, None, length=self.n_steps)
        return means

    @lazy_property
    def variance(self) -> Float[Array, "T M"]:
        """Marginal variance $\\mathrm{diag}(H P_t H^\\top + R)$, shape ``(T, M)``.

        The state recursion runs on dense ``P_t`` — it is the evolving
        quantity, so no input structure survives it — but the emission
        step avoids forming ``H P H^T``: operators go through
        `gaussx.sandwich` + `gaussx.diag`, arrays contract straight
        to the diagonal.
        """
        A_dense = _materialise(self.A)
        Q_dense = _materialise(self.Q)
        diag_R = _diag(_as_operator(self.R))
        H_op = self.H if isinstance(self.H, lx.AbstractLinearOperator) else None

        def step(P, _):
            P_next = A_dense @ P @ A_dense.T + Q_dense
            if H_op is not None:
                diag_HPH = _diag(sandwich(H_op, _psd(P_next)))
            else:
                diag_HPH = einsum(self.H, P_next, self.H, "m n, n k, m k -> m")
            return P_next, diag_HPH + diag_R

        _, variances = jax.lax.scan(step, self.P0, None, length=self.n_steps)
        return variances

n_steps: int property

Sequence length T.

mean() -> Float[Array, 'T M']

Marginal mean \(H A^t m_0\) of each observation, shape (T, M).

Source code in src/gaussx/_distributions/_lgssm.py

@lazy_property
def mean(self) -> Float[Array, "T M"]:
    """Marginal mean $H A^t m_0$ of each observation, shape ``(T, M)``."""

    def step(m, _):
        m_next = _matvec(self.A, m)
        return m_next, _matvec(self.H, m_next)

    _, means = jax.lax.scan(step, self.m0, None, length=self.n_steps)
    return means

variance() -> Float[Array, 'T M']

Marginal variance \(\mathrm{diag}(H P_t H^\top + R)\), shape (T, M).

The state recursion runs on dense P_t — it is the evolving quantity, so no input structure survives it — but the emission step avoids forming H P H^T: operators go through gaussx.sandwich + gaussx.diag, arrays contract straight to the diagonal.

Source code in src/gaussx/_distributions/_lgssm.py

@lazy_property
def variance(self) -> Float[Array, "T M"]:
    """Marginal variance $\\mathrm{diag}(H P_t H^\\top + R)$, shape ``(T, M)``.

    The state recursion runs on dense ``P_t`` — it is the evolving
    quantity, so no input structure survives it — but the emission
    step avoids forming ``H P H^T``: operators go through
    `gaussx.sandwich` + `gaussx.diag`, arrays contract straight
    to the diagonal.
    """
    A_dense = _materialise(self.A)
    Q_dense = _materialise(self.Q)
    diag_R = _diag(_as_operator(self.R))
    H_op = self.H if isinstance(self.H, lx.AbstractLinearOperator) else None

    def step(P, _):
        P_next = A_dense @ P @ A_dense.T + Q_dense
        if H_op is not None:
            diag_HPH = _diag(sandwich(H_op, _psd(P_next)))
        else:
            diag_HPH = einsum(self.H, P_next, self.H, "m n, n k, m k -> m")
        return P_next, diag_HPH + diag_R

    _, variances = jax.lax.scan(step, self.P0, None, length=self.n_steps)
    return variances

MaskedLGSSM

Bases: LGSSM

LGSSM whose log_prob marginalises unobserved channels exactly.

Takes a (T, M) boolean mask at construction. log_prob returns the exact marginal \(\log p(y_{\mathrm{obs}})\), not a bound: the conditional \(p(y_{\mathrm{miss}} \mid y_{\mathrm{obs}})\) is closed-form Gaussian, so dropping dimensions costs nothing in exactness. Masked entries of value are never read and may be NaN.

sample is inherited unchanged and simulates the complete (T, M) series — the mask governs which entries the density scores, not which are generated.

Parameters:

Name	Type	Description	Default
A	_Transition	Transition matrix or operator, shape (N, N).	required
H	_Emission	Emission matrix or operator, shape (M, N).	required
Q	_StateCov	Process noise covariance or operator, shape (N, N).	required
R	_ObsCov	Observation noise covariance or operator, shape (M, M).	required
m0	Float[Array, ' N']	Initial state mean, shape (N,).	required
P0	Float[Array, 'N N']	Initial state covariance, shape (N, N).	required
n_steps	int	Sequence length T.	required
obs_mask	Bool[Array, 'T M']	Observation mask, shape (T, M). True marks an observed channel. An all-False row leaves the state at its predicted value and contributes nothing. Named obs_mask so it does not shadow numpyro's inherited Distribution.mask() method.	required
validate_args	bool | None	Whether to validate input arguments.	None

Raises:

Type	Description
ValueError	If obs_mask does not have shape (n_steps, M).

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from gaussx import MaskedLGSSM
>>> mask = jr.bernoulli(jr.key(0), 0.7, (10, 2))
>>> d = MaskedLGSSM(0.9 * jnp.eye(2), jnp.eye(2), 0.1 * jnp.eye(2),
...                 0.2 * jnp.eye(2), jnp.zeros(2), jnp.eye(2),
...                 n_steps=10, obs_mask=mask)
>>> d.log_prob(d.sample(jr.key(1))).shape
()

Source code in src/gaussx/_distributions/_lgssm.py

class MaskedLGSSM(LGSSM):
    r"""`LGSSM` whose ``log_prob`` marginalises unobserved channels exactly.

    Takes a ``(T, M)`` boolean mask at construction. ``log_prob`` returns
    the exact marginal $\log p(y_{\mathrm{obs}})$, not a bound: the
    conditional $p(y_{\mathrm{miss}} \mid y_{\mathrm{obs}})$ is
    closed-form Gaussian, so dropping dimensions costs nothing in
    exactness. Masked entries of ``value`` are never read and may be
    ``NaN``.

    ``sample`` is inherited unchanged and simulates the **complete**
    ``(T, M)`` series — the mask governs which entries the density
    scores, not which are generated.

    Args:
        A: Transition matrix or operator, shape ``(N, N)``.
        H: Emission matrix or operator, shape ``(M, N)``.
        Q: Process noise covariance or operator, shape ``(N, N)``.
        R: Observation noise covariance or operator, shape ``(M, M)``.
        m0: Initial state mean, shape ``(N,)``.
        P0: Initial state covariance, shape ``(N, N)``.
        n_steps: Sequence length ``T``.
        obs_mask: Observation mask, shape ``(T, M)``. ``True`` marks an
            observed channel. An all-``False`` row leaves the state at
            its predicted value and contributes nothing. Named
            ``obs_mask`` so it does not shadow numpyro's inherited
            ``Distribution.mask()`` method.
        validate_args: Whether to validate input arguments.

    Raises:
        ValueError: If ``obs_mask`` does not have shape ``(n_steps, M)``.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from gaussx import MaskedLGSSM
        >>> mask = jr.bernoulli(jr.key(0), 0.7, (10, 2))
        >>> d = MaskedLGSSM(0.9 * jnp.eye(2), jnp.eye(2), 0.1 * jnp.eye(2),
        ...                 0.2 * jnp.eye(2), jnp.zeros(2), jnp.eye(2),
        ...                 n_steps=10, obs_mask=mask)
        >>> d.log_prob(d.sample(jr.key(1))).shape
        ()
    """

    pytree_data_fields = ("A", "H", "Q", "R", "m0", "P0", "obs_mask")

    # Named ``obs_mask`` rather than ``mask``: ``numpyro.distributions.
    # Distribution.mask`` is an inherited *method* returning a
    # ``MaskedDistribution`` (it backs ``numpyro.sample(..., obs_mask=)``),
    # and shadowing it with an array would silently break that API on
    # this subclass alone.
    obs_mask: Bool[Array, "T M"]

    def __init__(
        self,
        A: _Transition,
        H: _Emission,
        Q: _StateCov,
        R: _ObsCov,
        m0: Float[Array, " N"],
        P0: Float[Array, "N N"],
        n_steps: int,
        obs_mask: Bool[Array, "T M"],
        *,
        validate_args: bool | None = None,
    ) -> None:
        obs_mask = jnp.asarray(obs_mask, dtype=bool)
        expected = (n_steps, _out_size(H))
        if obs_mask.shape != expected:
            raise ValueError(
                f"obs_mask must have shape {expected}; got shape {obs_mask.shape}."
            )
        self.obs_mask = obs_mask
        super().__init__(A, H, Q, R, m0, P0, n_steps, validate_args=validate_args)

    @property
    def _obs_mask(self) -> Bool[Array, "T M"]:
        return self.obs_mask

    def _validate_sample(self, value: ArrayLike) -> ArrayLike:
        # Masked entries are never read and may be NaN, so check the support
        # on the observed entries only. numpyro >= 0.22 validates by default,
        # and a NaN anywhere would otherwise turn log_prob into -inf.
        return super()._validate_sample(jnp.where(self.obs_mask, value, 0.0))

LGSSMFactory

Bases: Module

Callable obs_mask -> MaskedLGSSM, for conditional use.

An equinox.Module rather than a closure so the state-space parameters stay visible to equinox.filter_grad when passed as gauss_flows.NumpyroBase(dist_factory=...).

Attributes:

Name	Type	Description
A	_Transition	Transition matrix or operator, shape (N, N).
H	_Emission	Emission matrix or operator, shape (M, N).
Q	_StateCov	Process noise covariance or operator, shape (N, N).
R	_ObsCov	Observation noise covariance or operator, shape (M, M).
m0	Float[Array, ' N']	Initial state mean, shape (N,).
P0	Float[Array, 'N N']	Initial state covariance, shape (N, N).
n_steps	int	Sequence length T.

Examples:

>>> import jax.numpy as jnp, jax.random as jr
>>> from gaussx import LGSSMFactory
>>> factory = LGSSMFactory(0.9 * jnp.eye(2), jnp.eye(2),
...                        0.1 * jnp.eye(2), 0.2 * jnp.eye(2),
...                        jnp.zeros(2), jnp.eye(2), n_steps=10)
>>> d = factory(jr.bernoulli(jr.key(0), 0.7, (10, 2)))
>>> d.event_shape
(10, 2)

Source code in src/gaussx/_distributions/_lgssm.py

class LGSSMFactory(eqx.Module):
    """Callable ``obs_mask -> MaskedLGSSM``, for conditional use.

    An `equinox.Module` rather than a closure so the state-space
    parameters stay visible to `equinox.filter_grad` when passed as
    ``gauss_flows.NumpyroBase(dist_factory=...)``.

    Attributes:
        A: Transition matrix or operator, shape ``(N, N)``.
        H: Emission matrix or operator, shape ``(M, N)``.
        Q: Process noise covariance or operator, shape ``(N, N)``.
        R: Observation noise covariance or operator, shape ``(M, M)``.
        m0: Initial state mean, shape ``(N,)``.
        P0: Initial state covariance, shape ``(N, N)``.
        n_steps: Sequence length ``T``.

    Examples:
        >>> import jax.numpy as jnp, jax.random as jr
        >>> from gaussx import LGSSMFactory
        >>> factory = LGSSMFactory(0.9 * jnp.eye(2), jnp.eye(2),
        ...                        0.1 * jnp.eye(2), 0.2 * jnp.eye(2),
        ...                        jnp.zeros(2), jnp.eye(2), n_steps=10)
        >>> d = factory(jr.bernoulli(jr.key(0), 0.7, (10, 2)))
        >>> d.event_shape
        (10, 2)
    """

    A: _Transition
    H: _Emission
    Q: _StateCov
    R: _ObsCov
    m0: Float[Array, " N"]
    P0: Float[Array, "N N"]
    n_steps: int = eqx.field(static=True)

    def __call__(self, obs_mask: Bool[Array, "T M"]) -> MaskedLGSSM:
        """Build the `MaskedLGSSM` for a given observation mask.

        Args:
            obs_mask: Observation mask, shape ``(T, M)``.

        Returns:
            A `MaskedLGSSM` carrying this factory's parameters.
        """
        return MaskedLGSSM(
            self.A,
            self.H,
            self.Q,
            self.R,
            self.m0,
            self.P0,
            self.n_steps,
            obs_mask,
        )

MarkovGaussian is the density over the states of a Gauss-Markov chain rather than its observations: \(x_0 \sim \mathcal{N}(\mu_0, P_0)\), \(x_{k+1} = A_k x_k + b_k + \varepsilon_k\), with event_shape \((T, d)\). sample and log_prob use the chain factorisation at \(O(T d^3)\) and never form the joint covariance, so it serves as a prior over a latent trajectory or as a structured variational guide. Its precision is block-tridiagonal; to_precision_form / from_precision_form convert to and from the \((\mu, \Lambda)\) layout that spingp_posterior returns, through the UDL factorisation, so a precision-form posterior becomes a sampleable chain in one \(O(T d^3)\) pass. Requires numpyro.

Structured linear algebra and Gaussian primitives for JAX.

MarkovGaussian

Bases: Distribution

Chain-Markov Gaussian over \(x_{0:T-1}\) as a density over (T, d).

The canonical parameterisation is the generative state-space form

\[ x_0 \sim \mathcal{N}(\mu_0, P_0), \qquad x_{k+1} = A_k x_k + b_k + \varepsilon_k, \quad \varepsilon_k \sim \mathcal{N}(0, Q_k), \]

whose joint covariance has a block-tridiagonal inverse. event_shape is (T, d), so log_prob of a trajectory is a scalar and the distribution can be used directly as a NumPyro site — a prior over a latent path, or a structured variational guide. Every view is derived from the SSM tuple on demand:

view	cost	route
marginals, mean, variance	\(O(T d^3)\)	forward moment propagation
sample, log_prob	\(O(T d^3)\)	chain factorisation, no joint Cholesky
precision, to_precision_form	\(O(T d^3)\)	udl_from_ssm_params
from_precision_form	\(O(T d^3)\)	udl_decomposition
covariance_matrix	\(O(T^2 d^3)\)	banded solves vs. identity; small \(T\)

Where LGSSM is a density over the observations of a chain with the states marginalised out, this is the density over the states themselves. Requires the numpyro optional extra (pip install "gaussx[numpyro]").

Parameters:

Name	Type	Description	Default
A	Float[Array, 'Tm1 d d']	Transition matrices, shape (T-1, d, d).	required
Q	Float[Array, 'Tm1 d d']	Process noise covariances, shape (T-1, d, d); Q[k] drives the step from \(x_k\) to \(x_{k+1}\).	required
mu0	Float[Array, ' d']	Initial mean, shape (d,).	required
P0	Float[Array, 'd d']	Initial covariance, shape (d, d).	required
b	Float[Array, 'Tm1 d'] | None	Optional transition offsets, shape (T-1, d). Defaults to zero.	None
validate_args	bool | None	Whether to validate input arguments.	None

Examples:

>>> import jax, jax.numpy as jnp, gaussx
>>> T, d = 10, 2
>>> A = jnp.broadcast_to(0.9 * jnp.eye(d), (T - 1, d, d))
>>> Q = jnp.broadcast_to(0.1 * jnp.eye(d), (T - 1, d, d))
>>> chain = gaussx.MarkovGaussian(A, Q, jnp.zeros(d), jnp.eye(d))
>>> xs = chain.sample(jax.random.key(0), (4,))
>>> xs.shape, chain.log_prob(xs).shape
((4, 10, 2), (4,))
>>> mean, precision = chain.to_precision_form()
>>> mean.shape, precision.as_matrix().shape
((20,), (20, 20))

Source code in src/gaussx/_distributions/_markov_gaussian.py

class MarkovGaussian(dist.Distribution):
    r"""Chain-Markov Gaussian over $x_{0:T-1}$ as a density over ``(T, d)``.

    The canonical parameterisation is the generative state-space form

    $$
    x_0 \sim \mathcal{N}(\mu_0, P_0), \qquad
    x_{k+1} = A_k x_k + b_k + \varepsilon_k, \quad
    \varepsilon_k \sim \mathcal{N}(0, Q_k),
    $$

    whose joint covariance has a block-tridiagonal inverse. ``event_shape``
    is ``(T, d)``, so ``log_prob`` of a trajectory is a scalar and the
    distribution can be used directly as a NumPyro site — a prior over a
    latent path, or a structured variational guide. Every view is derived
    from the SSM tuple on demand:

    | view | cost | route |
    |---|---|---|
    | `marginals`, `mean`, `variance` | $O(T d^3)$ | forward moment propagation |
    | `sample`, `log_prob` | $O(T d^3)$ | chain factorisation, no joint Cholesky |
    | `precision`, `to_precision_form` | $O(T d^3)$ | `udl_from_ssm_params` |
    | `from_precision_form` | $O(T d^3)$ | `udl_decomposition` |
    | `covariance_matrix` | $O(T^2 d^3)$ | banded solves vs. identity; small $T$ |

    Where `LGSSM` is a density over the *observations* of a chain with
    the states marginalised out, this is the density over the *states*
    themselves. Requires the ``numpyro`` optional extra
    (``pip install "gaussx[numpyro]"``).

    Args:
        A: Transition matrices, shape ``(T-1, d, d)``.
        Q: Process noise covariances, shape ``(T-1, d, d)``; ``Q[k]``
            drives the step from $x_k$ to $x_{k+1}$.
        mu0: Initial mean, shape ``(d,)``.
        P0: Initial covariance, shape ``(d, d)``.
        b: Optional transition offsets, shape ``(T-1, d)``. Defaults to
            zero.
        validate_args: Whether to validate input arguments.

    Examples:
        >>> import jax, jax.numpy as jnp, gaussx
        >>> T, d = 10, 2
        >>> A = jnp.broadcast_to(0.9 * jnp.eye(d), (T - 1, d, d))
        >>> Q = jnp.broadcast_to(0.1 * jnp.eye(d), (T - 1, d, d))
        >>> chain = gaussx.MarkovGaussian(A, Q, jnp.zeros(d), jnp.eye(d))
        >>> xs = chain.sample(jax.random.key(0), (4,))
        >>> xs.shape, chain.log_prob(xs).shape
        ((4, 10, 2), (4,))
        >>> mean, precision = chain.to_precision_form()
        >>> mean.shape, precision.as_matrix().shape
        ((20,), (20, 20))
    """

    arg_constraints = {}  # noqa: RUF012
    support = dist.constraints.real_matrix
    reparametrized_params = ["A", "b", "Q", "mu0", "P0"]  # noqa: RUF012
    pytree_data_fields = ("A", "b", "Q", "mu0", "P0")

    def __init__(
        self,
        A: Float[Array, "Tm1 d d"],
        Q: Float[Array, "Tm1 d d"],
        mu0: Float[Array, " d"],
        P0: Float[Array, "d d"],
        *,
        b: Float[Array, "Tm1 d"] | None = None,
        validate_args: bool | None = None,
    ) -> None:
        if A.ndim != 3 or A.shape[1] != A.shape[2]:
            raise ValueError(f"A must have shape (T-1, d, d), got {A.shape}.")
        Tm1, d, _ = A.shape
        if Q.shape != (Tm1, d, d):
            raise ValueError(f"Q must have shape {(Tm1, d, d)}, got {Q.shape}.")
        if mu0.shape != (d,):
            raise ValueError(f"mu0 must have shape {(d,)}, got {mu0.shape}.")
        if P0.shape != (d, d):
            raise ValueError(f"P0 must have shape {(d, d)}, got {P0.shape}.")
        if b is None:
            b = jnp.zeros((Tm1, d), dtype=A.dtype)
        elif b.shape != (Tm1, d):
            raise ValueError(f"b must have shape {(Tm1, d)}, got {b.shape}.")
        self.A = A
        self.b = b
        self.Q = Q
        self.mu0 = mu0
        self.P0 = P0
        super().__init__(
            batch_shape=(),
            event_shape=(Tm1 + 1, d),
            validate_args=validate_args,
        )

    @property
    def horizon(self) -> int:
        """Number of states ``T`` in the chain."""
        return self.event_shape[0]

    @property
    def state_dim(self) -> int:
        """State dimension ``d``."""
        return self.event_shape[1]

    # ------------------------------------------------------------------
    # Moment views
    # ------------------------------------------------------------------

    def marginals(self) -> tuple[Float[Array, "T d"], Float[Array, "T d d"]]:
        r"""Marginal means and covariances of every state.

        Forward propagation $m_{k+1} = A_k m_k + b_k$,
        $P_{k+1} = A_k P_k A_k^{\top} + Q_k$.

        Returns:
            Tuple ``(means, covs)`` of shapes ``(T, d)`` and ``(T, d, d)``.
        """

        def _step(carry, inputs):
            m, P = carry
            A_k, b_k, Q_k = inputs
            m_next = A_k @ m + b_k
            P_next = einsum(A_k, P, A_k, "i j, j k, l k -> i l") + Q_k
            return (m_next, P_next), (m_next, P_next)

        _, (means_rest, covs_rest) = jax.lax.scan(
            _step, (self.mu0, self.P0), (self.A, self.b, self.Q)
        )
        means = jnp.concatenate([self.mu0[None], means_rest], axis=0)
        covs = jnp.concatenate([self.P0[None], covs_rest], axis=0)
        return means, covs

    @lazy_property
    def mean(self) -> Float[Array, "T d"]:
        """Marginal means, shape ``(T, d)``."""
        return self.marginals()[0]

    @lazy_property
    def variance(self) -> Float[Array, "T d"]:
        """Marginal variances, shape ``(T, d)``."""
        _, covs = self.marginals()
        return jnp.diagonal(covs, axis1=-2, axis2=-1)

    def cross_covariances(self) -> Float[Array, "Tm1 d d"]:
        r"""$\mathrm{Cov}(x_{k+1}, x_k) = A_k P_k$ for each consecutive pair."""
        _, covs = self.marginals()
        return self._cross(covs)

    def _cross(self, covs: Float[Array, "T d d"]) -> Float[Array, "Tm1 d d"]:
        if self.A.shape[0] == 0:
            # One step, no transitions: einx cannot contract the empty
            # transition axis (gh-348).
            return self.A
        return einsum(self.A, covs[:-1], "T i j, T j k -> T i k")

    def pairwise_marginals(
        self,
    ) -> tuple[Float[Array, "Tm1 two_d"], Float[Array, "Tm1 two_d two_d"]]:
        r"""Joint $p(x_k, x_{k+1})$ for each consecutive pair.

        Returns:
            Tuple ``(joint_means, joint_covs)`` of shapes ``(T-1, 2d)``
            and ``(T-1, 2d, 2d)``, in the layout of
            `gaussx.pairwise_marginals`.
        """
        means, covs = self.marginals()
        if self.A.shape[0] == 0:
            d2 = 2 * self.state_dim
            return jnp.zeros((0, d2), means.dtype), jnp.zeros((0, d2, d2), covs.dtype)
        return pairwise_marginals(means, covs, self._cross(covs))

    @lazy_property
    def covariance_matrix(self) -> Float[Array, "Td Td"]:
        r"""Dense joint covariance $\Lambda^{-1}$, shape ``(Td, Td)``.

        Built column by column through `UDLDecomposition.solve`, so it
        costs $O(T^2 d^3)$ and is meant for small $T$ (reference checks,
        dense downstream consumers).
        """
        n = self.horizon * self.state_dim
        eye = jnp.eye(n, dtype=self.A.dtype)
        return jax.vmap(self.udl().solve, in_axes=1, out_axes=1)(eye)

    # ------------------------------------------------------------------
    # Precision views
    # ------------------------------------------------------------------

    def udl(self) -> UDLDecomposition:
        r"""The chain's precision, already factorised as $U \tilde{D} U^{\top}$."""
        return udl_from_ssm_params(self.A, self.Q, self.P0)

    @property
    def precision(self) -> BlockTriDiag:
        r"""Joint precision $\Lambda$ as a `BlockTriDiag`.

        Raw $\Lambda$, not the $-\tfrac12 \Lambda$ natural-parameter
        convention of `ssm_to_naturals`.
        """
        return self.udl().as_block_tridiag()

    def to_precision_form(self) -> tuple[Float[Array, " Td"], BlockTriDiag]:
        r"""The chain as ``(mean, precision)``.

        This is the layout `spingp_posterior` consumes and returns: pass
        ``mean`` as its ``prior_mean`` (it assumes a zero-mean prior
        otherwise) and hand its output to `from_precision_form`. The
        natural location parameter, if needed, is $\eta = \Lambda \mu$,
        i.e. ``precision.mv(mean)``.

        Returns:
            Tuple ``(mean, precision)``: the flattened marginal means of
            shape ``(T * d,)`` and the joint precision.
        """
        mean = rearrange(self.mean, "T d -> (T d)")
        return mean, self.precision

    @classmethod
    def from_precision_form(
        cls,
        mean: Float[Array, " Td"] | Float[Array, "T d"],
        precision: BlockTriDiag,
    ) -> MarkovGaussian:
        r"""Recover the generative chain from ``(mean, precision)``.

        One `udl_decomposition` pass gives $A_k = -U_k^{\top}$ and
        $Q_k = \tilde{D}_k^{-1}$; the offsets follow from the mean chain
        as $b_k = m_{k+1} - A_k m_k$. Round-trips `to_precision_form`
        exactly (to floating point).

        Args:
            mean: Marginal means, flattened ``(T * d,)`` or ``(T, d)``.
            precision: Symmetric positive-definite joint precision.

        Returns:
            The equivalent `MarkovGaussian`.
        """
        T, d = precision._num_blocks, precision._block_size
        means = mean if mean.ndim == 2 else rearrange(mean, "(T d) -> T d", T=T, d=d)
        A, Q, _ = udl_to_ssm_params(udl_decomposition(precision))
        if T == 1:
            b = means[1:]
        else:
            b = means[1:] - einsum(A, means[:-1], "T i j, T j -> T i")
        return cls(A, Q[1:], means[0], Q[0], b=b)

    # ------------------------------------------------------------------
    # Density and sampling
    # ------------------------------------------------------------------

    def _chol_factors(self) -> tuple[Float[Array, "d d"], Float[Array, "Tm1 d d"]]:
        return jnp.linalg.cholesky(self.P0), jnp.linalg.cholesky(self.Q)

    def _log_prob_single(self, xs: Float[Array, "T d"]) -> Float[Array, ""]:
        r"""$\log \mathcal{N}(x_0; \mu_0, P_0)
        + \sum_k \log \mathcal{N}(x_{k+1}; A_k x_k + b_k, Q_k)$ at $O(T d^3)$."""
        chol_P0, chol_Q = self._chol_factors()
        lp_0 = _chol_log_prob(chol_P0, xs[0] - self.mu0)
        if self.A.shape[0] == 0:
            return lp_0
        residuals = xs[1:] - einsum(self.A, xs[:-1], "T i j, T j -> T i") - self.b
        lp_rest = jax.vmap(_chol_log_prob)(chol_Q, residuals)
        return lp_0 + jnp.sum(lp_rest)

    @validate_sample
    def log_prob(self, value: Float[Array, "*batch T d"]) -> Float[Array, "*batch"]:
        leading_shape = value.shape[:-2]
        value_flat = rearrange(value, "... T d -> (...) T d")
        log_prob_flat = jax.vmap(self._log_prob_single)(value_flat)
        return _reshape_batch(log_prob_flat, leading_shape)

    def _sample_single(self, key: jax.Array) -> Float[Array, "T d"]:
        r"""One ancestral draw $x_0 \sim \mathcal{N}(\mu_0, P_0)$,
        $x_{k+1} = A_k x_k + b_k + L_k z_k$ at $O(T d^3)$."""
        chol_P0, chol_Q = self._chol_factors()
        T, d = self.horizon, self.state_dim
        k0, k_rest = jax.random.split(key)
        # Draw in the parameters' dtype, not JAX's default float.
        dtype = jnp.result_type(self.mu0, self.P0, self.Q)
        x0 = self.mu0 + chol_P0 @ jax.random.normal(k0, (d,), dtype=dtype)
        eps = jax.random.normal(k_rest, (T - 1, d), dtype=dtype)

        def _step(x, inputs):
            A_k, b_k, L_k, e_k = inputs
            x_next = A_k @ x + b_k + L_k @ e_k
            return x_next, x_next

        _, xs_rest = jax.lax.scan(_step, x0, (self.A, self.b, chol_Q, eps))
        return jnp.concatenate([x0[None], xs_rest], axis=0)

    def sample(
        self,
        key: jax.Array | None,
        sample_shape: tuple[int, ...] = (),
    ) -> Float[Array, "*sample T d"]:
        if key is None:
            raise ValueError("PRNG key must be provided to sample from MarkovGaussian.")
        n_samples = math.prod(sample_shape) if sample_shape else 1
        keys = jax.random.split(key, n_samples)
        samples_flat = jax.vmap(self._sample_single)(keys)
        return _reshape_series(samples_flat, sample_shape)

horizon: int property

Number of states T in the chain.

state_dim: int property

State dimension d.

precision: BlockTriDiag property

Joint precision \(\Lambda\) as a BlockTriDiag.

Raw \(\Lambda\), not the \(-\tfrac12 \Lambda\) natural-parameter convention of ssm_to_naturals.

marginals() -> tuple[Float[Array, 'T d'], Float[Array, 'T d d']]

Marginal means and covariances of every state.

Forward propagation \(m_{k+1} = A_k m_k + b_k\), \(P_{k+1} = A_k P_k A_k^{\top} + Q_k\).

Returns:

Type	Description
tuple[Float[Array, 'T d'], Float[Array, 'T d d']]	Tuple (means, covs) of shapes (T, d) and (T, d, d).

Source code in src/gaussx/_distributions/_markov_gaussian.py

def marginals(self) -> tuple[Float[Array, "T d"], Float[Array, "T d d"]]:
    r"""Marginal means and covariances of every state.

    Forward propagation $m_{k+1} = A_k m_k + b_k$,
    $P_{k+1} = A_k P_k A_k^{\top} + Q_k$.

    Returns:
        Tuple ``(means, covs)`` of shapes ``(T, d)`` and ``(T, d, d)``.
    """

    def _step(carry, inputs):
        m, P = carry
        A_k, b_k, Q_k = inputs
        m_next = A_k @ m + b_k
        P_next = einsum(A_k, P, A_k, "i j, j k, l k -> i l") + Q_k
        return (m_next, P_next), (m_next, P_next)

    _, (means_rest, covs_rest) = jax.lax.scan(
        _step, (self.mu0, self.P0), (self.A, self.b, self.Q)
    )
    means = jnp.concatenate([self.mu0[None], means_rest], axis=0)
    covs = jnp.concatenate([self.P0[None], covs_rest], axis=0)
    return means, covs

mean() -> Float[Array, 'T d']

Marginal means, shape (T, d).

Source code in src/gaussx/_distributions/_markov_gaussian.py

@lazy_property
def mean(self) -> Float[Array, "T d"]:
    """Marginal means, shape ``(T, d)``."""
    return self.marginals()[0]

variance() -> Float[Array, 'T d']

Marginal variances, shape (T, d).

Source code in src/gaussx/_distributions/_markov_gaussian.py

@lazy_property
def variance(self) -> Float[Array, "T d"]:
    """Marginal variances, shape ``(T, d)``."""
    _, covs = self.marginals()
    return jnp.diagonal(covs, axis1=-2, axis2=-1)

cross_covariances() -> Float[Array, 'Tm1 d d']

\(\mathrm{Cov}(x_{k+1}, x_k) = A_k P_k\) for each consecutive pair.

Source code in src/gaussx/_distributions/_markov_gaussian.py

def cross_covariances(self) -> Float[Array, "Tm1 d d"]:
    r"""$\mathrm{Cov}(x_{k+1}, x_k) = A_k P_k$ for each consecutive pair."""
    _, covs = self.marginals()
    return self._cross(covs)

pairwise_marginals() -> tuple[Float[Array, 'Tm1 two_d'], Float[Array, 'Tm1 two_d two_d']]

Joint \(p(x_k, x_{k+1})\) for each consecutive pair.

Returns:

Type	Description
Float[Array, 'Tm1 two_d']	Tuple (joint_means, joint_covs) of shapes (T-1, 2d)
Float[Array, 'Tm1 two_d two_d']	and (T-1, 2d, 2d), in the layout of
tuple[Float[Array, 'Tm1 two_d'], Float[Array, 'Tm1 two_d two_d']]	gaussx.pairwise_marginals.

Source code in src/gaussx/_distributions/_markov_gaussian.py

def pairwise_marginals(
    self,
) -> tuple[Float[Array, "Tm1 two_d"], Float[Array, "Tm1 two_d two_d"]]:
    r"""Joint $p(x_k, x_{k+1})$ for each consecutive pair.

    Returns:
        Tuple ``(joint_means, joint_covs)`` of shapes ``(T-1, 2d)``
        and ``(T-1, 2d, 2d)``, in the layout of
        `gaussx.pairwise_marginals`.
    """
    means, covs = self.marginals()
    if self.A.shape[0] == 0:
        d2 = 2 * self.state_dim
        return jnp.zeros((0, d2), means.dtype), jnp.zeros((0, d2, d2), covs.dtype)
    return pairwise_marginals(means, covs, self._cross(covs))

covariance_matrix() -> Float[Array, 'Td Td']

Dense joint covariance \(\Lambda^{-1}\), shape (Td, Td).

Built column by column through UDLDecomposition.solve, so it costs \(O(T^2 d^3)\) and is meant for small \(T\) (reference checks, dense downstream consumers).

Source code in src/gaussx/_distributions/_markov_gaussian.py

@lazy_property
def covariance_matrix(self) -> Float[Array, "Td Td"]:
    r"""Dense joint covariance $\Lambda^{-1}$, shape ``(Td, Td)``.

    Built column by column through `UDLDecomposition.solve`, so it
    costs $O(T^2 d^3)$ and is meant for small $T$ (reference checks,
    dense downstream consumers).
    """
    n = self.horizon * self.state_dim
    eye = jnp.eye(n, dtype=self.A.dtype)
    return jax.vmap(self.udl().solve, in_axes=1, out_axes=1)(eye)

udl() -> UDLDecomposition

The chain's precision, already factorised as \(U \tilde{D} U^{\top}\).

Source code in src/gaussx/_distributions/_markov_gaussian.py

def udl(self) -> UDLDecomposition:
    r"""The chain's precision, already factorised as $U \tilde{D} U^{\top}$."""
    return udl_from_ssm_params(self.A, self.Q, self.P0)

to_precision_form() -> tuple[Float[Array, ' Td'], BlockTriDiag]

The chain as (mean, precision).

This is the layout spingp_posterior consumes and returns: pass mean as its prior_mean (it assumes a zero-mean prior otherwise) and hand its output to from_precision_form. The natural location parameter, if needed, is \(\eta = \Lambda \mu\), i.e. precision.mv(mean).

Returns:

Type	Description
Float[Array, ' Td']	Tuple (mean, precision): the flattened marginal means of
BlockTriDiag	shape (T * d,) and the joint precision.

Source code in src/gaussx/_distributions/_markov_gaussian.py

def to_precision_form(self) -> tuple[Float[Array, " Td"], BlockTriDiag]:
    r"""The chain as ``(mean, precision)``.

    This is the layout `spingp_posterior` consumes and returns: pass
    ``mean`` as its ``prior_mean`` (it assumes a zero-mean prior
    otherwise) and hand its output to `from_precision_form`. The
    natural location parameter, if needed, is $\eta = \Lambda \mu$,
    i.e. ``precision.mv(mean)``.

    Returns:
        Tuple ``(mean, precision)``: the flattened marginal means of
        shape ``(T * d,)`` and the joint precision.
    """
    mean = rearrange(self.mean, "T d -> (T d)")
    return mean, self.precision

from_precision_form(mean: Float[Array, ' Td'] | Float[Array, 'T d'], precision: BlockTriDiag) -> MarkovGaussian classmethod

Recover the generative chain from (mean, precision).

One udl_decomposition pass gives \(A_k = -U_k^{\top}\) and \(Q_k = \tilde{D}_k^{-1}\); the offsets follow from the mean chain as \(b_k = m_{k+1} - A_k m_k\). Round-trips to_precision_form exactly (to floating point).

Parameters:

Name	Type	Description	Default
mean	Float[Array, ' Td'] | Float[Array, 'T d']	Marginal means, flattened (T * d,) or (T, d).	required
precision	BlockTriDiag	Symmetric positive-definite joint precision.	required

Returns:

Type	Description
MarkovGaussian	The equivalent MarkovGaussian.

Source code in src/gaussx/_distributions/_markov_gaussian.py

@classmethod
def from_precision_form(
    cls,
    mean: Float[Array, " Td"] | Float[Array, "T d"],
    precision: BlockTriDiag,
) -> MarkovGaussian:
    r"""Recover the generative chain from ``(mean, precision)``.

    One `udl_decomposition` pass gives $A_k = -U_k^{\top}$ and
    $Q_k = \tilde{D}_k^{-1}$; the offsets follow from the mean chain
    as $b_k = m_{k+1} - A_k m_k$. Round-trips `to_precision_form`
    exactly (to floating point).

    Args:
        mean: Marginal means, flattened ``(T * d,)`` or ``(T, d)``.
        precision: Symmetric positive-definite joint precision.

    Returns:
        The equivalent `MarkovGaussian`.
    """
    T, d = precision._num_blocks, precision._block_size
    means = mean if mean.ndim == 2 else rearrange(mean, "(T d) -> T d", T=T, d=d)
    A, Q, _ = udl_to_ssm_params(udl_decomposition(precision))
    if T == 1:
        b = means[1:]
    else:
        b = means[1:] - einsum(A, means[:-1], "T i j, T j -> T i")
    return cls(A, Q[1:], means[0], Q[0], b=b)

Gaussian sugar ops

\[ \log \mathcal{N}(x \mid \mu, \Sigma) = -\tfrac12 (x-\mu)^\top \Sigma^{-1} (x-\mu) - \tfrac12 \log|\Sigma| - \tfrac{N}{2}\log 2\pi \]

evaluated through structured solve + logdet, plus entropy, quadratic forms, KL divergences, conditioning, and the numerically stable Joseph-form covariance update.

Structured linear algebra and Gaussian primitives for JAX.

gaussian_log_prob(loc: Float[Array, ' N'], cov_operator: lx.AbstractLinearOperator, value: Float[Array, ' N'], *, solver: AbstractSolverStrategy | None = None, key: jax.Array | None = None) -> Float[Array, '']

Multivariate normal log-probability.

Computes:

log N(value | loc, Sigma)
= -0.5 * (N log(2 pi) + log|Sigma| + (value - loc)^T Sigma^{-1} (value - loc))

All expensive operations (solve, logdet) dispatch on operator structure automatically, or through an explicit solver.

Parameters:

Name	Type	Description	Default
loc	Float[Array, ' N']	Mean vector, shape (N,).	required
cov_operator	AbstractLinearOperator	Covariance operator Sigma, shape (N, N).	required
value	Float[Array, ' N']	Observation vector, shape (N,).	required
solver	AbstractSolverStrategy | None	Optional solver strategy (needs both solve and logdet). When None, uses structural dispatch.	None
key	Array | None	PRNG key for a stochastic logdet strategy's probes. None uses the strategy's own seed, i.e. the same probes on every call (see gaussx.KeyedSolver). Ignored by exact strategies.	None

Returns:

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

Source code in src/gaussx/_distributions/_gaussian.py

def gaussian_log_prob(
    loc: Float[Array, " N"],
    cov_operator: lx.AbstractLinearOperator,
    value: Float[Array, " N"],
    *,
    solver: AbstractSolverStrategy | None = None,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Multivariate normal log-probability.

    Computes:

        log N(value | loc, Sigma)
        = -0.5 * (N log(2 pi) + log|Sigma| + (value - loc)^T Sigma^{-1} (value - loc))

    All expensive operations (``solve``, ``logdet``) dispatch on
    operator structure automatically, or through an explicit *solver*.

    Args:
        loc: Mean vector, shape ``(N,)``.
        cov_operator: Covariance operator Sigma, shape ``(N, N)``.
        value: Observation vector, shape ``(N,)``.
        solver: Optional solver strategy (needs both solve and logdet).
            When ``None``, uses structural dispatch.
        key: PRNG key for a stochastic logdet strategy's probes. ``None``
            uses the strategy's own seed, i.e. the same probes on every call
            (see `gaussx.KeyedSolver`). Ignored by exact strategies.

    Returns:
        Scalar log-probability.
    """
    return _gaussian_log_prob_residual(
        value - loc, cov_operator, solver=solver, key=key
    )

gaussian_entropy(cov_operator: lx.AbstractLinearOperator, *, solver: AbstractLogdetStrategy | None = None, key: jax.Array | None = None) -> Float[Array, '']

Entropy of a multivariate normal N(mu, Sigma).

Computes:

H = 0.5 * (N * (1 + log(2 pi)) + log|Sigma|)

Independent of the mean.

Parameters:

Name	Type	Description	Default
cov_operator	AbstractLinearOperator	Covariance operator, shape (N, N).	required
solver	AbstractLogdetStrategy | None	Optional logdet strategy. When None, uses structural dispatch.	None
key	Array | None	PRNG key for a stochastic logdet strategy's probes. None uses the strategy's own seed, i.e. the same probes on every call (see gaussx.KeyedSolver). Ignored by exact strategies.	None

Returns:

Type	Description
Float[Array, '']	Scalar entropy.

Source code in src/gaussx/_distributions/_gaussian.py

def gaussian_entropy(
    cov_operator: lx.AbstractLinearOperator,
    *,
    solver: AbstractLogdetStrategy | None = None,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """Entropy of a multivariate normal ``N(mu, Sigma)``.

    Computes:

        H = 0.5 * (N * (1 + log(2 pi)) + log|Sigma|)

    Independent of the mean.

    Args:
        cov_operator: Covariance operator, shape ``(N, N)``.
        solver: Optional logdet strategy. When ``None``, uses
            structural dispatch.
        key: PRNG key for a stochastic logdet strategy's probes. ``None``
            uses the strategy's own seed, i.e. the same probes on every call
            (see `gaussx.KeyedSolver`). Ignored by exact strategies.

    Returns:
        Scalar entropy.
    """
    N = cov_operator.in_size()
    ld = dispatch_logdet(cov_operator, solver, key=key)
    return 0.5 * (N * (1.0 + _LOG_2PI) + ld)

quadratic_form(operator: lx.AbstractLinearOperator, x: Float[Array, ' N'], *, solver: AbstractSolveStrategy | None = None) -> Float[Array, '']

Compute x^T A^{-1} x via a single solve.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A non-singular linear operator A.	required
x	Float[Array, ' N']	Vector, shape (N,).	required
solver	AbstractSolveStrategy | None	Optional solve strategy. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, '']	Scalar x^T A^{-1} x.

Source code in src/gaussx/_distributions/_gaussian.py

def quadratic_form(
    operator: lx.AbstractLinearOperator,
    x: Float[Array, " N"],
    *,
    solver: AbstractSolveStrategy | None = None,
) -> Float[Array, ""]:
    """Compute ``x^T A^{-1} x`` via a single solve.

    Args:
        operator: A non-singular linear operator A.
        x: Vector, shape ``(N,)``.
        solver: Optional solve strategy. When ``None``, uses
            structural dispatch.

    Returns:
        Scalar ``x^T A^{-1} x``.
    """
    return x @ dispatch_solve(operator, x, solver)

kl_standard_normal(m: Float[Array, ' N'], S: lx.AbstractLinearOperator, *, solver: AbstractLogdetStrategy | None = None, key: jax.Array | None = None) -> Float[Array, '']

KL(first ‖ second): KL(N(m, S) || N(0, I)).

Special case of gaussian_kl with q_loc = 0 and q_cov = I. The identity prior means no matrix inversion is required, making this more efficient than calling the general form directly.

Computes:

KL = 0.5 * (tr(S) + m^T m - N - log|S|)

Ubiquitous in variational inference as the prior KL term.

Parameters:

Name	Type	Description	Default
m	Float[Array, ' N']	Mean vector, shape (N,).	required
S	AbstractLinearOperator	Covariance operator, shape (N, N).	required
solver	AbstractLogdetStrategy | None	Optional logdet strategy. When None, uses structural dispatch.	None
key	Array | None	PRNG key for a stochastic logdet strategy's probes. None uses the strategy's own seed, i.e. the same probes on every call (see gaussx.KeyedSolver). Ignored by exact strategies.	None

Returns:

Type	Description
Float[Array, '']	Scalar KL divergence.

See Also

gaussian_kl: General KL between two multivariate normals with arbitrary lineax covariance operators.

Source code in src/gaussx/_distributions/_gaussian.py

def kl_standard_normal(
    m: Float[Array, " N"],
    S: lx.AbstractLinearOperator,
    *,
    solver: AbstractLogdetStrategy | None = None,
    key: jax.Array | None = None,
) -> Float[Array, ""]:
    """KL(first ‖ second): ``KL(N(m, S) || N(0, I))``.

    Special case of `gaussian_kl`
    with ``q_loc = 0`` and ``q_cov = I``.  The identity prior means no
    matrix inversion is required, making this more efficient than calling
    the general form directly.

    Computes:

        KL = 0.5 * (tr(S) + m^T m - N - log|S|)

    Ubiquitous in variational inference as the prior KL term.

    Args:
        m: Mean vector, shape ``(N,)``.
        S: Covariance operator, shape ``(N, N)``.
        solver: Optional logdet strategy. When ``None``, uses
            structural dispatch.
        key: PRNG key for a stochastic logdet strategy's probes. ``None``
            uses the strategy's own seed, i.e. the same probes on every call
            (see `gaussx.KeyedSolver`). Ignored by exact strategies.

    Returns:
        Scalar KL divergence.

    See Also:
        `gaussian_kl`: General KL
        between two multivariate normals with arbitrary lineax covariance
        operators.
    """
    N = m.shape[-1]
    tr_S = trace(S)
    mTm = m @ m
    ld = dispatch_logdet(S, solver, key=key)
    return 0.5 * (tr_S + mTm - N - ld)

gaussian_kl(p_loc: Float[Array, ' N'], p_cov: lx.AbstractLinearOperator, q_loc: Float[Array, ' N'], q_cov: lx.AbstractLinearOperator) -> Float[Array, '']

KL(first ‖ second): KL(p || q) between two multivariate normals.

This is the canonical KL implementation for lineax-operator covariances. The specialised variants below all compute the same quantity but with different parameterisations suited to their use cases:

• kl_standard_normal — special case KL(N(m, S) || N(0, I)); avoids matrix inversion.
• gauss_kl — Cholesky-parameterised form for GP/SVGP models; supports multi-output and diagonal q_sqrt.
• kl_divergence — Bregman-divergence form operating on natural parameters for the exponential family.

\[ KL(p \| q) = \frac{1}{2}\bigl( \operatorname{tr}(\Sigma_q^{-1} \Sigma_p) + (\mu_q - \mu_p)^T \Sigma_q^{-1} (\mu_q - \mu_p) - N + \log|\Sigma_q| - \log|\Sigma_p| \bigr) \]

Exploits structured operators for the trace and logdet terms.

Parameters:

Name	Type	Description	Default
p_loc	Float[Array, ' N']	Mean of distribution p, shape (N,).	required
p_cov	AbstractLinearOperator	Covariance operator of distribution p.	required
q_loc	Float[Array, ' N']	Mean of distribution q, shape (N,).	required
q_cov	AbstractLinearOperator	Covariance operator of distribution q.	required

Returns:

Type	Description
Float[Array, '']	Scalar KL divergence.

Source code in src/gaussx/_distributions/_kl.py

def gaussian_kl(
    p_loc: Float[Array, " N"],
    p_cov: lx.AbstractLinearOperator,
    q_loc: Float[Array, " N"],
    q_cov: lx.AbstractLinearOperator,
) -> Float[Array, ""]:
    r"""KL(first ‖ second): ``KL(p || q)`` between two multivariate normals.

    This is the **canonical KL implementation** for lineax-operator covariances.
    The specialised variants below all compute the same quantity but with
    different parameterisations suited to their use cases:

    - `kl_standard_normal` —
      special case ``KL(N(m, S) || N(0, I))``; avoids matrix inversion.
    - `gauss_kl` — Cholesky-parameterised form
      for GP/SVGP models; supports multi-output and diagonal ``q_sqrt``.
    - `kl_divergence` — Bregman-divergence
      form operating on natural parameters for the exponential family.

    $$
    KL(p \| q) = \frac{1}{2}\bigl(
        \operatorname{tr}(\Sigma_q^{-1} \Sigma_p)
        + (\mu_q - \mu_p)^T \Sigma_q^{-1} (\mu_q - \mu_p)
        - N
        + \log|\Sigma_q| - \log|\Sigma_p|
    \bigr)
    $$

    Exploits structured operators for the trace and logdet terms.

    Args:
        p_loc: Mean of distribution p, shape ``(N,)``.
        p_cov: Covariance operator of distribution p.
        q_loc: Mean of distribution q, shape ``(N,)``.
        q_cov: Covariance operator of distribution q.

    Returns:
        Scalar KL divergence.
    """
    N = p_loc.shape[-1]
    delta = q_loc - p_loc

    p_matrix, q_matrix = _dense_psd_matrix(p_cov), _dense_psd_matrix(q_cov)
    if p_matrix is not None and q_matrix is not None:
        # Two Choleskys for the whole KL (gh-329): with M = L_q^{-1} L_p,
        # tr(Sigma_q^{-1} Sigma_p) = ||M||_F^2, and the quadratic form and
        # both log-determinants come from the same factors.
        L_p = jnp.linalg.cholesky(p_matrix)
        L_q = jnp.linalg.cholesky(q_matrix)
        M = jax.scipy.linalg.solve_triangular(L_q, L_p, lower=True)
        z = jax.scipy.linalg.solve_triangular(L_q, delta, lower=True)
        ld_diff = _cholesky_logdet(L_q) - _cholesky_logdet(L_p)
        return 0.5 * (jnp.sum(M**2) + z @ z - N + ld_diff)

    # tr(Sigma_q^{-1} Sigma_p)
    q_inv = inv(q_cov)
    trace_term = trace_product(q_inv, p_cov)

    # Quadratic term: delta^T Sigma_q^{-1} delta
    quad = jnp.sum(delta * solve(q_cov, delta))

    # Log-determinant difference
    ld_q = logdet(q_cov)
    ld_p = logdet(p_cov)

    return 0.5 * (trace_term + quad - N + ld_q - ld_p)

conditional(loc: Float[Array, ' N'], cov: lx.AbstractLinearOperator, obs_idx: Int[Array, ' M'], obs_values: Float[Array, ' M'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' R'], lx.AbstractLinearOperator]

Compute p(x_A | x_B = b) from a joint Gaussian p(x_A, x_B).

Given a joint distribution \(\mathcal{N}(\mu, \Sigma)\) and observed indices B with values b, returns the conditional distribution over the remaining indices A:

\[ \begin{aligned} \mu_{A|B} &= \mu_A + \Sigma_{AB} \Sigma_{BB}^{-1} (b - \mu_B) \\ \Sigma_{A|B} &= \Sigma_{AA} - \Sigma_{AB} \Sigma_{BB}^{-1} \Sigma_{BA} \end{aligned} \]

Parameters:

Name	Type	Description	Default
loc	Float[Array, ' N']	Mean vector of the joint distribution, shape (N,).	required
cov	AbstractLinearOperator	Covariance operator of the joint distribution, shape (N, N).	required
obs_idx	Int[Array, ' M']	Indices of the observed variables, shape (M,).	required
obs_values	Float[Array, ' M']	Observed values, shape (M,).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

Type	Description
Float[Array, ' R']	Tuple (cond_mean, cond_cov) — mean and covariance of the
AbstractLinearOperator	conditional distribution over unobserved variables.

Raises:

Type	Description
ValueError	If obs_idx is not 1D, does not match obs_values in shape, or -- when it is concrete -- is out of bounds or contains duplicates.
EquinoxRuntimeError	At run time, if a traced obs_idx is out of bounds or contains duplicates.

Note

Usable under jax.jit. Concrete indices (NumPy or JAX constants, including ones closed over by the jitted function) are validated on the host; traced indices are validated at run time with equinox.error_if, since an out-of-bounds gather would otherwise be silently clamped.

Source code in src/gaussx/_distributions/_conditional.py

def conditional(
    loc: Float[Array, " N"],
    cov: lx.AbstractLinearOperator,
    obs_idx: Int[Array, " M"],
    obs_values: Float[Array, " M"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " R"], lx.AbstractLinearOperator]:
    r"""Compute ``p(x_A | x_B = b)`` from a joint Gaussian ``p(x_A, x_B)``.

    Given a joint distribution $\mathcal{N}(\mu, \Sigma)$ and
    observed indices *B* with values *b*, returns the conditional
    distribution over the remaining indices *A*:

    $$
    \begin{aligned}
    \mu_{A|B} &= \mu_A + \Sigma_{AB} \Sigma_{BB}^{-1} (b - \mu_B) \\
    \Sigma_{A|B} &= \Sigma_{AA} - \Sigma_{AB} \Sigma_{BB}^{-1} \Sigma_{BA}
    \end{aligned}
    $$

    Args:
        loc: Mean vector of the joint distribution, shape ``(N,)``.
        cov: Covariance operator of the joint distribution, shape ``(N, N)``.
        obs_idx: Indices of the observed variables, shape ``(M,)``.
        obs_values: Observed values, shape ``(M,)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(cond_mean, cond_cov)`` — mean and covariance of the
        conditional distribution over unobserved variables.

    Raises:
        ValueError: If ``obs_idx`` is not 1D, does not match
            ``obs_values`` in shape, or -- when it is concrete -- is out
            of bounds or contains duplicates.
        EquinoxRuntimeError: At run time, if a *traced* ``obs_idx`` is
            out of bounds or contains duplicates.

    Note:
        Usable under ``jax.jit``. Concrete indices (NumPy or JAX constants,
        including ones closed over by the jitted function) are validated
        on the host; traced indices are validated at run time with
        `equinox.error_if`, since an out-of-bounds gather would otherwise
        be silently clamped.
    """
    N = loc.shape[0]
    obs_values = jnp.asarray(obs_values, dtype=loc.dtype)

    if isinstance(obs_idx, Tracer):
        obs_idx = obs_idx.astype(jnp.int32)
        _check_shapes(obs_idx, obs_values)
        bad = jnp.any((obs_idx < 0) | (obs_idx >= N)) | jnp.any(
            jnp.diff(jnp.sort(obs_idx)) == 0
        )
        # Attached to obs_idx, which feeds both outputs, so the check cannot
        # be dead-code-eliminated when a caller uses only one of them.
        obs_idx = eqx.error_if(
            obs_idx,
            bad,
            f"obs_idx must be within bounds [0, {N}) and must not contain duplicates.",
        )
    else:
        # Cast before validating, so e.g. [1.1, 1.2] is caught as the
        # duplicate [1, 1] it becomes.
        idx_np = np.asarray(obs_idx).astype(np.int32)
        _check_shapes(idx_np, obs_values)
        if np.any((idx_np < 0) | (idx_np >= N)):
            raise ValueError(f"obs_idx must be within bounds [0, {N}).")
        if np.any(np.diff(np.sort(idx_np)) == 0):
            raise ValueError("obs_idx must not contain duplicates.")
        obs_idx = jnp.asarray(idx_np)

    # Build mask for unobserved indices
    mask = jnp.ones(N, dtype=bool).at[obs_idx].set(False)
    free_idx = jnp.where(mask, size=N - obs_idx.shape[0])[0]

    # Extract sub-blocks via structural dispatch — avoids materializing
    # the full ``(N, N)`` joint covariance for structured operators
    # (Diagonal, BlockDiag). Falls back to full materialization for
    # unstructured operators where there is no efficient alternative.
    Sigma_AA = submatrix(cov, free_idx, free_idx)
    Sigma_AB = submatrix(cov, free_idx, obs_idx)
    Sigma_BB = submatrix(cov, obs_idx, obs_idx)

    mu_A = loc[free_idx]
    mu_B = loc[obs_idx]

    # Sigma_BB^{-1} (b - mu_B)
    residual = obs_values - mu_B
    Sigma_BB_op = lx.MatrixLinearOperator(Sigma_BB, lx.positive_semidefinite_tag)
    alpha = dispatch_solve(Sigma_BB_op, residual, solver)

    # Conditional mean: mu_A + Sigma_AB @ alpha
    cond_mean = mu_A + Sigma_AB @ alpha

    # Sigma_BB^{-1} Sigma_BA — single matrix solve (one factorization for
    # the whole RHS in the default PSD path).
    Sigma_BA = Sigma_AB.T
    X = solve_matrix(Sigma_BB_op, Sigma_BA, solver=solver)

    # Conditional covariance: Sigma_AA - Sigma_AB @ X
    cond_cov_mat = Sigma_AA - Sigma_AB @ X

    # Symmetrize for numerical stability
    cond_cov_mat = symmetrize(cond_cov_mat)
    cond_cov = lx.MatrixLinearOperator(cond_cov_mat, lx.positive_semidefinite_tag)

    return cond_mean, cond_cov

joseph_update(P_pred: Float[Array, 'N N'], K: Float[Array, 'N M'], H: Float[Array, 'M N'], R: Float[Array, 'M M']) -> Float[Array, 'N N']

Numerically stable Joseph-form covariance update.

Computes the updated covariance after a Kalman measurement update:

P_update = (I - K H) P_pred (I - K H)^T + K R K^T

This form is more numerically stable than the simplified P = P_pred - K S K^T or P = (I - K H) P_pred because it guarantees symmetry and is more robust when the Kalman gain K is approximate or the system is poorly conditioned.

Parameters:

Name	Type	Description	Default
P_pred	Float[Array, 'N N']	Predicted covariance, shape (N, N).	required
K	Float[Array, 'N M']	Kalman gain, shape (N, M).	required
H	Float[Array, 'M N']	Observation model, shape (M, N).	required
R	Float[Array, 'M M']	Observation noise covariance, shape (M, M).	required

Returns:

Type	Description
Float[Array, 'N N']	Updated covariance, shape (N, N).

Source code in src/gaussx/_distributions/_joseph.py

def joseph_update(
    P_pred: Float[Array, "N N"],
    K: Float[Array, "N M"],
    H: Float[Array, "M N"],
    R: Float[Array, "M M"],
) -> Float[Array, "N N"]:
    r"""Numerically stable Joseph-form covariance update.

    Computes the updated covariance after a Kalman measurement update:

        P_update = (I - K H) P_pred (I - K H)^T + K R K^T

    This form is more numerically stable than the simplified
    ``P = P_pred - K S K^T`` or ``P = (I - K H) P_pred`` because it
    guarantees symmetry and is more robust when the Kalman gain ``K``
    is approximate or the system is poorly conditioned.

    Args:
        P_pred: Predicted covariance, shape ``(N, N)``.
        K: Kalman gain, shape ``(N, M)``.
        H: Observation model, shape ``(M, N)``.
        R: Observation noise covariance, shape ``(M, M)``.

    Returns:
        Updated covariance, shape ``(N, N)``.
    """
    N = P_pred.shape[0]
    I_KH = jnp.eye(N, dtype=P_pred.dtype) - K @ H  # (N, N)
    P_update = I_KH @ P_pred @ I_KH.T + K @ R @ K.T
    return (P_update + P_update.T) / 2

add_jitter(operator: lx.AbstractLinearOperator, jitter: float = 1e-06) -> lx.AbstractLinearOperator

Add diagonal jitter for numerical stability: A + eps * I.

Parameters:

Name	Type	Description	Default
operator	AbstractLinearOperator	A linear operator, shape (N, N).	required
jitter	float	Scalar jitter value. Default 1e-6.	1e-06

Returns:

Type	Description
AbstractLinearOperator	A + jitter * I as a lineax AddLinearOperator.

Source code in src/gaussx/_distributions/_gaussian.py

def add_jitter(
    operator: lx.AbstractLinearOperator,
    jitter: float = 1e-6,
) -> lx.AbstractLinearOperator:
    """Add diagonal jitter for numerical stability: ``A + eps * I``.

    Args:
        operator: A linear operator, shape ``(N, N)``.
        jitter: Scalar jitter value. Default ``1e-6``.

    Returns:
        ``A + jitter * I`` as a lineax ``AddLinearOperator``.
    """
    n = operator.in_size()
    dtype = operator.out_structure().dtype
    jitter_op = lx.DiagonalLinearOperator(jnp.full(n, jitter, dtype=dtype))
    return operator + jitter_op

project(K_XZ: Float[Array, 'B M'], L_Z: lx.AbstractLinearOperator) -> Float[Array, 'B M']

Compute A_X = K_XZ @ K_ZZ^{-1} via Cholesky solve.

Solves L_Z @ L_Z^T @ A_X^T = K_XZ^T using forward/backward substitution. Used in sparse variational GPs to project test points onto the inducing space.

Parameters:

Name	Type	Description	Default
K_XZ	Float[Array, 'B M']	Cross-covariance matrix, shape (B, M).	required
L_Z	AbstractLinearOperator	Lower-triangular Cholesky factor of K_ZZ, shape (M, M). It need not carry lineax.lower_triangular_tag: the tag is added when missing (the factor is lower-triangular by precondition), and L_Z.T is then upper-triangular.	required

Returns:

Type	Description
Float[Array, 'B M']	Projection matrix A_X, shape (B, M).

Raises:

Type	Description
ValueError	If L_Z is not (M, M) with M = K_XZ.shape[1].

Source code in src/gaussx/_distributions/_project.py

def project(
    K_XZ: Float[Array, "B M"],
    L_Z: lx.AbstractLinearOperator,
) -> Float[Array, "B M"]:
    """Compute A_X = K_XZ @ K_ZZ^{-1} via Cholesky solve.

    Solves ``L_Z @ L_Z^T @ A_X^T = K_XZ^T`` using forward/backward
    substitution.  Used in sparse variational GPs to project test
    points onto the inducing space.

    Args:
        K_XZ: Cross-covariance matrix, shape ``(B, M)``.
        L_Z: Lower-triangular Cholesky factor of K_ZZ, shape ``(M, M)``.
            It need not carry ``lineax.lower_triangular_tag``: the tag is
            added when missing (the factor is lower-triangular by
            precondition), and ``L_Z.T`` is then upper-triangular.

    Returns:
        Projection matrix A_X, shape ``(B, M)``.

    Raises:
        ValueError: If ``L_Z`` is not ``(M, M)`` with ``M = K_XZ.shape[1]``.
    """
    M = K_XZ.shape[-1]
    if L_Z.in_size() != M or L_Z.out_size() != M:
        raise ValueError(
            f"L_Z must be ({M}, {M}) to match K_XZ's {M} columns, got "
            f"({L_Z.out_size()}, {L_Z.in_size()})."
        )
    if not lx.is_lower_triangular(L_Z):
        # lx.Triangular checks tags, so an untagged factor (e.g. built from
        # jnp.linalg.cholesky) would be rejected.
        L_Z = lx.TaggedLinearOperator(L_Z, lx.lower_triangular_tag)
    # Solve L_Z @ Y = K_XZ^T, then L_Z^T @ A_X^T = Y
    # Equivalently, solve (L_Z @ L_Z^T) @ A_X^T = K_XZ^T per column
    solver = lx.Triangular()

    def _solve_col(kxz_row):
        # Solve L_Z y = kxz_row
        y = lx.linear_solve(L_Z, kxz_row, solver).value
        # Solve L_Z^T a = y
        return lx.linear_solve(L_Z.T, y, solver).value

    return jax.vmap(_solve_col)(K_XZ)

Structured sampling

sample_mvn draws \(x = \mu + L\varepsilon\) with \(LL^\top = K\), picking the factor \(L\) from the structure of the covariance operator: element-wise for a diagonal, per factor for a Kronecker product, per block for a block-diagonal, FFT circulant embedding for a Toeplitz, and \(L_B z_1 + U\sqrt{D}z_2\) for a low-rank update, so a \(10^4 \times 10^4\) Kronecker covariance is never factorised densely. Every branch is exact. A SumOfKroneckers therefore takes the dense symmetric square root: its matrix-free sampler, sum_of_kroneckers_sample, is a truncated Lanczos approximation and stays an explicit opt-in. The one-off helpers toeplitz_sample and kronecker_sum_sample also remain available for callers that want them directly.

sample_joint_conditional draws a partitioned Gaussian \((x_a, x_b)\) jointly without forming its \((N + M)^2\) covariance: the observed block through sample_mvn, so it keeps its structure, and the other through the gain \(K_{to}K_{oo}^{-1}\) plus a square root of the Schur complement. Given an observed value it also returns conditional draws that share the same noise, which makes them exactly the matheron_update of the joint draws.

Structured linear algebra and Gaussian primitives for JAX.

sample_mvn(mean: Float[Array, '*batch N'], covariance: lx.AbstractLinearOperator, *, key: jax.Array, num_samples: int = 1) -> Float[Array, 'num_samples *batch N']

Draw samples from \(\mathcal{N}(\mu, K)\), dispatching on \(K\)'s structure.

Every branch draws \(x = \mu + L\varepsilon\) with \(\varepsilon \sim \mathcal{N}(0, I)\) and \(LL^{\top} = K\) exactly -- no branch approximates the covariance -- and only the factor \(L\) changes:

Covariance \(K\)	Factor	Cost
lineax.DiagonalLinearOperator	\(\sqrt{d} \odot z\)	\(O(N)\)
gaussx.BlockDiag	each block by its own rule	per block
gaussx.Kronecker	per-factor eigendecomposition	\(O(\sum_i N_i^3)\)
gaussx.KroneckerSum	per-factor eigendecomposition	\(O(n_A^3 + n_B^3)\)
gaussx.BlockTriDiag	block-banded Cholesky	\(O(T d^3)\)
gaussx.Toeplitz	FFT circulant embedding	\(O(N \log N)\)
gaussx.LowRankUpdate	\(L_B z_1 + U\sqrt{D} z_2\)	\(B\)'s cost \(+ O(Nr)\)
anything else	dense symmetric square root	\(O(N^3)\)

lineax.TaggedLinearOperator wrappers are looked through, and a scalar multiple \(cK\) (with \(c > 0\)) samples \(K\) and rescales by \(\sqrt{c}\), so wrapping never costs the structure underneath.

The Kronecker branches work in the factors' eigenbases, where the eigenvalues of the whole operator are products (or sums) of the factors' ones. That keeps them exact when a factor is indefinite but the product is not -- (-A) \otimes (-B) -- and when the covariance is only semi-definite. The dense fallback is the symmetric square root for the same reason: a Cholesky factor of a singular covariance is NaN.

A gaussx.SumOfKroneckers takes the dense fallback: its matrix-free square root, gaussx.sum_of_kroneckers_sample, is a truncated Lanczos approximation, so it stays an explicit opt-in rather than a default.

Some structured routes hold only for some values, and each takes the dense square root otherwise -- so every covariance the operator can represent is sampled exactly, and pathwise gradients stay finite:

• a scalar multiple \(cK\) needs \(c > 0\);
• a gaussx.Toeplitz needs a circulant embedding (2, 4, 8 or 16 times N) whose spectrum is strictly positive -- a zero eigenvalue there would put sqrt(0) in the pathwise derivative;
• a gaussx.LowRankUpdate \(B + UDV^{\top}\) needs \(V = U\) established (the same array, or equal values), a base \(B\) that is a covariance in its own right (tagged positive semi-definite, or a strictly positive diagonal), and \(D > 0\) -- a zero weight would again put sqrt(0) in the derivative.

When those values are only known at run time (under jax.jit), the choice is a jax.lax.cond, which runs only the branch it takes; under jax.vmap it evaluates both.

Parameters:

Name	Type	Description	Default
mean	Float[Array, '*batch N']	Mean \(\mu\) with shape (*batch, N). Each batch element gets independent noise; the covariance is shared.	required
covariance	AbstractLinearOperator	Positive-definite covariance operator \(K\) of size (N, N).	required
key	Array	PRNG key.	required
num_samples	int	Number of independent draws.	1

Returns:

Type	Description
Float[Array, 'num_samples *batch N']	Samples of shape (num_samples, *batch, N).

Raises:

Type	Description
ValueError	If covariance is not square, does not match the last axis of mean, or num_samples is below one.

Source code in src/gaussx/_distributions/_sample.py

def sample_mvn(
    mean: Float[Array, "*batch N"],
    covariance: lx.AbstractLinearOperator,
    *,
    key: jax.Array,
    num_samples: int = 1,
) -> Float[Array, "num_samples *batch N"]:
    r"""Draw samples from $\mathcal{N}(\mu, K)$, dispatching on $K$'s structure.

    Every branch draws $x = \mu + L\varepsilon$ with $\varepsilon \sim
    \mathcal{N}(0, I)$ and $LL^{\top} = K$ *exactly* -- no branch
    approximates the covariance -- and only the factor $L$ changes:

    | Covariance $K$ | Factor | Cost |
    |---|---|---|
    | `lineax.DiagonalLinearOperator` | $\sqrt{d} \odot z$ | $O(N)$ |
    | `gaussx.BlockDiag` | each block by its own rule | per block |
    | `gaussx.Kronecker` | per-factor eigendecomposition | $O(\sum_i N_i^3)$ |
    | `gaussx.KroneckerSum` | per-factor eigendecomposition | $O(n_A^3 + n_B^3)$ |
    | `gaussx.BlockTriDiag` | block-banded Cholesky | $O(T d^3)$ |
    | `gaussx.Toeplitz` | FFT circulant embedding | $O(N \log N)$ |
    | `gaussx.LowRankUpdate` | $L_B z_1 + U\sqrt{D} z_2$ | $B$'s cost $+ O(Nr)$ |
    | anything else | dense symmetric square root | $O(N^3)$ |

    `lineax.TaggedLinearOperator` wrappers are looked through, and a scalar
    multiple $cK$ (with $c > 0$) samples $K$ and rescales by $\sqrt{c}$, so
    wrapping never costs the structure underneath.

    The Kronecker branches work in the factors' eigenbases, where the
    eigenvalues of the whole operator are products (or sums) of the factors'
    ones. That keeps them exact when a factor is indefinite but the product
    is not -- ``(-A) \otimes (-B)`` -- and when the covariance is only
    semi-definite. The dense fallback is the symmetric square root for the
    same reason: a Cholesky factor of a singular covariance is ``NaN``.

    A `gaussx.SumOfKroneckers` takes the dense fallback: its matrix-free
    square root, `gaussx.sum_of_kroneckers_sample`, is a truncated Lanczos
    approximation, so it stays an explicit opt-in rather than a default.

    Some structured routes hold only for some values, and each takes the
    dense square root otherwise -- so every covariance the operator can
    represent is sampled exactly, and pathwise gradients stay finite:

    - a scalar multiple $cK$ needs $c > 0$;
    - a `gaussx.Toeplitz` needs a circulant embedding (2, 4, 8 or 16 times
      ``N``) whose spectrum is strictly positive -- a zero eigenvalue there
      would put ``sqrt(0)`` in the pathwise derivative;
    - a `gaussx.LowRankUpdate` $B + UDV^{\top}$ needs $V = U$ established
      (the same array, or equal values), a base $B$ that is a covariance in
      its own right (tagged positive semi-definite, or a strictly positive
      diagonal), and $D > 0$ -- a zero weight would again put ``sqrt(0)``
      in the derivative.

    When those values are only known at run time (under `jax.jit`), the
    choice is a `jax.lax.cond`, which runs only the branch it takes; under
    `jax.vmap` it evaluates both.

    Args:
        mean: Mean $\mu$ with shape ``(*batch, N)``. Each batch element gets
            independent noise; the covariance is shared.
        covariance: Positive-definite covariance operator $K$ of size
            ``(N, N)``.
        key: PRNG key.
        num_samples: Number of independent draws.

    Returns:
        Samples of shape ``(num_samples, *batch, N)``.

    Raises:
        ValueError: If ``covariance`` is not square, does not match the last
            axis of ``mean``, or ``num_samples`` is below one.
    """
    mean = jnp.asarray(mean)
    if covariance.in_size() != covariance.out_size():
        raise ValueError(
            "sample_mvn requires a square covariance, got "
            f"{covariance.out_size()}x{covariance.in_size()}."
        )
    if mean.ndim < 1 or mean.shape[-1] != covariance.in_size():
        raise ValueError(
            f"mean must have shape (*batch, {covariance.in_size()}), got {mean.shape}."
        )
    if num_samples < 1:
        raise ValueError(f"num_samples must be at least 1, got {num_samples}.")

    batch_shape = mean.shape[:-1]
    num_draws = num_samples * math.prod(batch_shape)
    if num_draws == 0:
        # An empty batch: nothing to draw, but the output keeps its shape.
        # The represented dtype, which `in_structure` can understate (a
        # LowRankUpdate reports its base's). Abstract, so nothing is formed.
        matrix_dtype = jax.eval_shape(covariance.as_matrix).dtype
        dtype = jnp.result_type(mean, matrix_dtype)
        return jnp.zeros((num_samples, *mean.shape), dtype=dtype)
    draws = _zero_mean_draws(covariance, key, num_draws)

    axes = [f"b{index}" for index in range(len(batch_shape))]
    pattern = f"(s {' '.join(axes)}) n -> s {' '.join(axes)} n"
    draws = rearrange(draws, pattern, **dict(zip(axes, batch_shape, strict=True)))
    dtype = jnp.result_type(mean, draws)
    return mean.astype(dtype) + draws.astype(dtype)

sample_joint_conditional(joint_mean: tuple[Float[Array, ' N'], Float[Array, ' M']], joint_covariance: dict[str, lx.AbstractLinearOperator], *, key: jax.Array, observed_index: int = 1, observed_value: Float[Array, ' K'] | None = None, num_samples: int = 1, solver: AbstractSolveStrategy | None = None) -> dict[str, Any]

Joint draws of a partitioned Gaussian, and conditional draws from them.

For

\[ \begin{pmatrix} x_a \\ x_b \end{pmatrix} \sim \mathcal{N}\!\left( \begin{pmatrix} m_a \\ m_b \end{pmatrix}, \begin{pmatrix} K_{aa} & K_{ab} \\ K_{ba} & K_{bb} \end{pmatrix} \right), \]

write \(o\) for the observed block and \(t\) for the other (with observed_index=1, \(o = b\) and \(t = a\)). The joint draw uses the block factorisation

\[ x_o = m_o + L_o z_1, \qquad x_t = m_t + K_{to} K_{oo}^{-1} (x_o - m_o) + L_S z_2, \]

with \(L_o L_o^{\top} = K_{oo}\) drawn by gaussx.sample_mvn -- so the observed block keeps its structure -- and \(L_S L_S^{\top} = S = K_{tt} - K_{to} K_{oo}^{-1} K_{ot}\) the Schur complement. The joint covariance of size \((N + M)^2\) is never formed.

When observed_value \(\beta\) is given, the conditional draws of \(x_t \mid x_o = \beta\) reuse the same \(z_2\):

\[ x_t \mid \beta = m_t + K_{to} K_{oo}^{-1} (\beta - m_o) + L_S z_2, \]

which is exactly gaussx.matheron_update applied to the joint draws.

The Schur complement is dense in general: \(K_{to} K_{oo}^{-1} K_{ot}\) is dense whatever the structure of \(K_{tt}\), so \(S\) is factorised at \(O(N_t^3)\) plus \(N_t\) solves against \(K_{oo}\). Its square root is the symmetric one from dense_symmetric_sqrt, with round-off negative eigenvalues clipped. That tolerates the rank deficiency of target points that coincide with observed ones -- a Cholesky would return NaNs there -- and its custom JVP keeps gradients finite when \(S\) has repeated eigenvalues.

Parameters:

Name	Type	Description	Default
joint_mean	tuple[Float[Array, ' N'], Float[Array, ' M']]	(m_a, m_b) with shapes (N,) and (M,).	required
joint_covariance	dict[str, AbstractLinearOperator]	Operators {"aa": K_aa, "ab": K_ab, "bb": K_bb} of shapes (N, N), (N, M) and (M, M); K_ba is K_ab.T.	required
key	Array	PRNG key.	required
observed_index	int	Which block is conditioned on: 1 for \(x_b\) (sample \(x_a \mid x_b\)), 0 for \(x_a\).	1
observed_value	Float[Array, ' K'] | None	Optional value \(\beta\) of the observed block. When given, the result also holds conditional draws.	None
num_samples	int	Number of independent draws.	1
solver	AbstractSolveStrategy | None	Optional solver strategy for the solves against \(K_{oo}\) (e.g. gaussx.CGSolver). When None, routes through structural dispatch, as gaussx.matheron_update does.	None

Returns:

Type	Description
dict[str, Any]	{"joint": (samples_a, samples_b)} with shapes
dict[str, Any]	(num_samples, N) and (num_samples, M), plus
dict[str, Any]	"conditional" of shape (num_samples, N_t) when
dict[str, Any]	observed_value is given.

Raises:

Type	Description
ValueError	On a missing covariance block, mismatched shapes, observed_index outside {0, 1} or num_samples below one.

Source code in src/gaussx/_distributions/_sample.py

def sample_joint_conditional(
    joint_mean: tuple[Float[Array, " N"], Float[Array, " M"]],
    joint_covariance: dict[str, lx.AbstractLinearOperator],
    *,
    key: jax.Array,
    observed_index: int = 1,
    observed_value: Float[Array, " K"] | None = None,
    num_samples: int = 1,
    solver: AbstractSolveStrategy | None = None,
) -> dict[str, Any]:
    r"""Joint draws of a partitioned Gaussian, and conditional draws from them.

    For

    $$
    \begin{pmatrix} x_a \\ x_b \end{pmatrix} \sim \mathcal{N}\!\left(
    \begin{pmatrix} m_a \\ m_b \end{pmatrix},
    \begin{pmatrix} K_{aa} & K_{ab} \\ K_{ba} & K_{bb} \end{pmatrix}
    \right),
    $$

    write $o$ for the observed block and $t$ for the other (with
    ``observed_index=1``, $o = b$ and $t = a$). The joint draw uses the block
    factorisation

    $$
    x_o = m_o + L_o z_1, \qquad
    x_t = m_t + K_{to} K_{oo}^{-1} (x_o - m_o) + L_S z_2,
    $$

    with $L_o L_o^{\top} = K_{oo}$ drawn by `gaussx.sample_mvn` -- so the
    observed block keeps its structure -- and $L_S L_S^{\top} = S = K_{tt} -
    K_{to} K_{oo}^{-1} K_{ot}$ the Schur complement. The joint covariance of
    size $(N + M)^2$ is never formed.

    When ``observed_value`` $\beta$ is given, the conditional draws of
    $x_t \mid x_o = \beta$ reuse the same $z_2$:

    $$
    x_t \mid \beta = m_t + K_{to} K_{oo}^{-1} (\beta - m_o) + L_S z_2,
    $$

    which is exactly `gaussx.matheron_update` applied to the joint draws.

    The Schur complement is dense in general: $K_{to} K_{oo}^{-1} K_{ot}$ is
    dense whatever the structure of $K_{tt}$, so $S$ is factorised at
    $O(N_t^3)$ plus $N_t$ solves against $K_{oo}$. Its square root is the
    symmetric one from ``dense_symmetric_sqrt``, with round-off negative
    eigenvalues clipped. That tolerates the rank deficiency of target points
    that coincide with observed ones -- a Cholesky would return NaNs there --
    and its custom JVP keeps gradients finite when $S$ has repeated
    eigenvalues.

    Args:
        joint_mean: ``(m_a, m_b)`` with shapes ``(N,)`` and ``(M,)``.
        joint_covariance: Operators ``{"aa": K_aa, "ab": K_ab, "bb": K_bb}``
            of shapes ``(N, N)``, ``(N, M)`` and ``(M, M)``; ``K_ba`` is
            ``K_ab.T``.
        key: PRNG key.
        observed_index: Which block is conditioned on: ``1`` for $x_b$
            (sample $x_a \mid x_b$), ``0`` for $x_a$.
        observed_value: Optional value $\beta$ of the observed block. When
            given, the result also holds conditional draws.
        num_samples: Number of independent draws.
        solver: Optional solver strategy for the solves against $K_{oo}$
            (e.g. `gaussx.CGSolver`). When ``None``, routes through
            structural dispatch, as `gaussx.matheron_update` does.

    Returns:
        ``{"joint": (samples_a, samples_b)}`` with shapes
        ``(num_samples, N)`` and ``(num_samples, M)``, plus
        ``"conditional"`` of shape ``(num_samples, N_t)`` when
        ``observed_value`` is given.

    Raises:
        ValueError: On a missing covariance block, mismatched shapes,
            ``observed_index`` outside ``{0, 1}`` or ``num_samples`` below
            one.
    """
    mean_a, mean_b = (jnp.asarray(mean) for mean in joint_mean)
    covariance_aa, covariance_ab, covariance_bb = _covariance_blocks(
        joint_covariance, mean_a.shape, mean_b.shape
    )
    if observed_index not in (0, 1):
        raise ValueError(f"observed_index must be 0 or 1, got {observed_index}.")
    if num_samples < 1:
        raise ValueError(f"num_samples must be at least 1, got {num_samples}.")

    if observed_index == 1:
        mean_o, mean_t = mean_b, mean_a
        covariance_oo, covariance_tt = covariance_bb, covariance_aa
        cross = covariance_ab
    else:
        mean_o, mean_t = mean_a, mean_b
        covariance_oo, covariance_tt = covariance_aa, covariance_bb
        cross = covariance_ab.T

    if observed_value is not None:
        observed_value = jnp.asarray(observed_value)
        if observed_value.shape != mean_o.shape:
            raise ValueError(
                f"observed_value must have shape {mean_o.shape}, got "
                f"{observed_value.shape}."
            )

    if solver is None:
        solve_observed = lambda vector: solve(covariance_oo, vector)
    else:
        solve_observed = lambda vector: dispatch_solve(covariance_oo, vector, solver)

    def gain(deviations: Float[Array, "S K"]) -> Float[Array, "S T"]:
        """Apply ``K_to K_oo⁺`` to each row.

        The structured solve is used whenever it succeeds. A singular -- only
        semi-definite -- ``K_oo`` makes it return non-finite values, and then
        the batch is redone with the pseudo-inverse: for a valid joint
        covariance both the deviations and the rows of ``K_to`` lie in the
        range of ``K_oo``, where the pseudo-inverse is exact. ``lax.cond``
        evaluates only the branch it takes.
        """
        solved = jax.vmap(solve_observed)(deviations)
        solved = jax.lax.cond(
            jnp.all(jnp.isfinite(solved)),
            lambda: solved,
            lambda: _pseudo_inverse_solve(covariance_oo, deviations),
        )
        return jax.vmap(cross.mv)(solved)

    observed_key, schur_key = jr.split(key)
    deviations_o = sample_mvn(
        jnp.zeros_like(mean_o), covariance_oo, key=observed_key, num_samples=num_samples
    )
    dtype = deviations_o.dtype
    schur_root = _schur_root(covariance_tt, cross, gain)
    noise = jr.normal(schur_key, (num_samples, mean_t.shape[0]), dtype=dtype)
    residual = einsum(noise, schur_root, "s k, t k -> s t")

    samples_o = mean_o + deviations_o
    samples_t = mean_t + gain(deviations_o) + residual
    joint = (samples_t, samples_o) if observed_index == 1 else (samples_o, samples_t)
    result: dict[str, Any] = {"joint": joint}
    if observed_value is not None:
        # Matheron's update of the returned joint draws, literally: with a
        # finite-tolerance solver the gain is not additive across right-hand
        # sides, so shifting from the mean instead would drift from it.
        result["conditional"] = samples_t + gain(observed_value - samples_o)
    return result

Exponential family

The Gaussian in natural form: \(\eta_1 = \Lambda\mu\), \(\eta_2 = -\tfrac12 \Lambda\). Conversions between mean/covariance, natural, and expectation parameterizations — multivariate (operator-aware) and univariate (per-site diagonal) — plus the log-partition, Fisher information, and sufficient statistics that natural-gradient and EP updates are built from.

Structured linear algebra and Gaussian primitives for JAX.

GaussianExpFam

Bases: Module

Gaussian in natural (exponential family) parameters.

\[ q(x \mid \eta) = h(x) \exp(\eta^T T(x) - A(\eta)) \]

where:

• Natural parameters: eta1 = Lambda @ mu, eta2 = -0.5 * Lambda
• Sufficient statistics: T(x) = [x, x x^T]
• Log-partition (as log_partition returns it): ``A(eta) = -0.25 * eta1^T eta2^{-1} eta1 - 0.5 * log|-2 eta2|
• 0.5 * N * log(2 pi)``
• Base measure: h(x) = 1, so log q(x) = eta^T T(x) - A(eta). The (2 pi)^{-N/2} normaliser lives in A, not in h.

Attributes:

Name	Type	Description
eta1	Float[Array, ' N']	Natural location parameter, shape (N,).
eta2	AbstractLinearOperator	Natural precision-like operator, shape (N, N). Represents -0.5 * Lambda where Lambda is the precision.

Source code in src/gaussx/_expfam/_gaussian.py

class GaussianExpFam(eqx.Module):
    r"""Gaussian in natural (exponential family) parameters.

    $$
    q(x \mid \eta) = h(x) \exp(\eta^T T(x) - A(\eta))
    $$

    where:

    - Natural parameters: ``eta1 = Lambda @ mu``, ``eta2 = -0.5 * Lambda``
    - Sufficient statistics: ``T(x) = [x, x x^T]``
    - Log-partition (as `log_partition` returns it):
      ``A(eta) = -0.25 * eta1^T eta2^{-1} eta1 - 0.5 * log|-2 eta2|
      + 0.5 * N * log(2 pi)``
    - Base measure: ``h(x) = 1``, so ``log q(x) = eta^T T(x) - A(eta)``.
      The ``(2 pi)^{-N/2}`` normaliser lives in ``A``, not in ``h``.

    Attributes:
        eta1: Natural location parameter, shape ``(N,)``.
        eta2: Natural precision-like operator, shape ``(N, N)``.
            Represents ``-0.5 * Lambda`` where Lambda is the precision.
    """

    eta1: Float[Array, " N"]
    eta2: lx.AbstractLinearOperator

    @staticmethod
    def from_mean_cov(
        mu: Float[Array, " N"],
        Sigma: lx.AbstractLinearOperator,
    ) -> GaussianExpFam:
        """Construct from mean and covariance.

        Args:
            mu: Mean vector, shape ``(N,)``.
            Sigma: Covariance operator, shape ``(N, N)``.

        Returns:
            A ``GaussianExpFam`` instance.
        """
        eta1, eta2 = mean_cov_to_natural(mu, Sigma)
        return GaussianExpFam(eta1=eta1, eta2=eta2)

    @staticmethod
    def from_mean_prec(
        mu: Float[Array, " N"],
        Lambda: lx.AbstractLinearOperator,
    ) -> GaussianExpFam:
        """Construct from mean and precision.

        Args:
            mu: Mean vector, shape ``(N,)``.
            Lambda: Precision operator, shape ``(N, N)``.

        Returns:
            A ``GaussianExpFam`` instance.
        """
        eta1 = Lambda.mv(mu)
        eta2 = -0.5 * Lambda
        return GaussianExpFam(eta1=eta1, eta2=eta2)

from_mean_cov(mu: Float[Array, ' N'], Sigma: lx.AbstractLinearOperator) -> GaussianExpFam staticmethod

Construct from mean and covariance.

Parameters:

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

Returns:

Type	Description
GaussianExpFam	A GaussianExpFam instance.

Source code in src/gaussx/_expfam/_gaussian.py

@staticmethod
def from_mean_cov(
    mu: Float[Array, " N"],
    Sigma: lx.AbstractLinearOperator,
) -> GaussianExpFam:
    """Construct from mean and covariance.

    Args:
        mu: Mean vector, shape ``(N,)``.
        Sigma: Covariance operator, shape ``(N, N)``.

    Returns:
        A ``GaussianExpFam`` instance.
    """
    eta1, eta2 = mean_cov_to_natural(mu, Sigma)
    return GaussianExpFam(eta1=eta1, eta2=eta2)

from_mean_prec(mu: Float[Array, ' N'], Lambda: lx.AbstractLinearOperator) -> GaussianExpFam staticmethod

Construct from mean and precision.

Parameters:

Name	Type	Description	Default
mu	Float[Array, ' N']	Mean vector, shape (N,).	required
Lambda	AbstractLinearOperator	Precision operator, shape (N, N).	required

Returns:

Type	Description
GaussianExpFam	A GaussianExpFam instance.

Source code in src/gaussx/_expfam/_gaussian.py

@staticmethod
def from_mean_prec(
    mu: Float[Array, " N"],
    Lambda: lx.AbstractLinearOperator,
) -> GaussianExpFam:
    """Construct from mean and precision.

    Args:
        mu: Mean vector, shape ``(N,)``.
        Lambda: Precision operator, shape ``(N, N)``.

    Returns:
        A ``GaussianExpFam`` instance.
    """
    eta1 = Lambda.mv(mu)
    eta2 = -0.5 * Lambda
    return GaussianExpFam(eta1=eta1, eta2=eta2)

to_mean_cov(expfam: GaussianExpFam) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator]

Convert natural parameters to the mean and covariance.

The inverse of GaussianExpFam.from_mean_cov. These are not the expectation parameters (mu, mu mu^T + Sigma): for those, use gaussx.natural_to_expectation.

Parameters:

Name	Type	Description	Default
expfam	GaussianExpFam	Gaussian in natural form.	required

Returns:

Type	Description
tuple[Float[Array, ' N'], AbstractLinearOperator]	Tuple (mu, Sigma) — mean vector and covariance operator.

Source code in src/gaussx/_expfam/_gaussian.py

def to_mean_cov(
    expfam: GaussianExpFam,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Convert natural parameters to the mean and covariance.

    The inverse of `GaussianExpFam.from_mean_cov`. These are **not** the
    expectation parameters ``(mu, mu mu^T + Sigma)``: for those, use
    `gaussx.natural_to_expectation`.

    Args:
        expfam: Gaussian in natural form.

    Returns:
        Tuple ``(mu, Sigma)`` — mean vector and covariance operator.
    """
    return natural_to_mean_cov(expfam.eta1, expfam.eta2)

mean_cov_to_natural(mu: Float[Array, ' N'], Sigma: lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator]

Convert mean/covariance to natural parameters (operator form).

Given mean mu and covariance Sigma:

• eta1 = solve(Sigma, mu)
• eta2 = -0.5 * inv(Sigma)

Operator structure (diagonal, Kronecker, …) is exploited via structural dispatch. For dense-array inputs see mean_chol_to_natural.

For block-tridiagonal (SSM) inputs see gaussx._ssm._ssm_natural.ssm_to_naturals.

Parameters:

Name	Type	Description	Default
mu	Float[Array, ' N']	Mean vector, shape (N,).	required
Sigma	AbstractLinearOperator	Covariance operator, shape (N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, ' N']	Tuple (eta1, eta2) where eta1 is shape (N,) and
AbstractLinearOperator	eta2 is a linear operator.

Source code in src/gaussx/_expfam/_natural.py

def mean_cov_to_natural(
    mu: Float[Array, " N"],
    Sigma: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Convert mean/covariance to natural parameters (operator form).

    Given mean ``mu`` and covariance ``Sigma``:

    - ``eta1 = solve(Sigma, mu)``
    - ``eta2 = -0.5 * inv(Sigma)``

    Operator structure (diagonal, Kronecker, …) is exploited via
    structural dispatch. For dense-array inputs see
    `mean_chol_to_natural`.

    For block-tridiagonal (SSM) inputs see
    `gaussx._ssm._ssm_natural.ssm_to_naturals`.

    Args:
        mu: Mean vector, shape ``(N,)``.
        Sigma: Covariance operator, shape ``(N, N)``.
        solver: Optional solver strategy. When ``None``, uses
            structural dispatch.

    Returns:
        Tuple ``(eta1, eta2)`` where eta1 is shape ``(N,)`` and
        eta2 is a linear operator.
    """
    eta1 = dispatch_solve(Sigma, mu, solver)
    eta2 = -0.5 * inv(Sigma)
    return eta1, eta2

natural_to_mean_cov(eta1: Float[Array, ' N'], eta2: lx.AbstractLinearOperator, *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator]

Convert natural parameters to mean/covariance (operator form).

Given natural parameters (eta1, eta2) where eta1 = Lambda @ mu and eta2 = -0.5 * Lambda:

• mu = solve(-2 * eta2, eta1)
• Sigma = inv(-2 * eta2)

Operator structure (diagonal, Kronecker, …) is exploited via structural dispatch. For dense-array inputs see natural_to_mean_chol.

For block-tridiagonal (SSM) inputs see gaussx._ssm._ssm_natural.naturals_to_ssm.

Parameters:

Name	Type	Description	Default
eta1	Float[Array, ' N']	Natural location parameter, shape (N,).	required
eta2	AbstractLinearOperator	Natural precision-like operator, shape (N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy. When None, uses structural dispatch.	None

Returns:

Type	Description
Float[Array, ' N']	Tuple (mu, Sigma) where mu is shape (N,) and
AbstractLinearOperator	Sigma is a linear operator.

Source code in src/gaussx/_expfam/_natural.py

def natural_to_mean_cov(
    eta1: Float[Array, " N"],
    eta2: lx.AbstractLinearOperator,
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Convert natural parameters to mean/covariance (operator form).

    Given natural parameters ``(eta1, eta2)`` where
    ``eta1 = Lambda @ mu`` and ``eta2 = -0.5 * Lambda``:

    - ``mu = solve(-2 * eta2, eta1)``
    - ``Sigma = inv(-2 * eta2)``

    Operator structure (diagonal, Kronecker, …) is exploited via
    structural dispatch. For dense-array inputs see
    `natural_to_mean_chol`.

    For block-tridiagonal (SSM) inputs see
    `gaussx._ssm._ssm_natural.naturals_to_ssm`.

    Args:
        eta1: Natural location parameter, shape ``(N,)``.
        eta2: Natural precision-like operator, shape ``(N, N)``.
        solver: Optional solver strategy. When ``None``, uses
            structural dispatch.

    Returns:
        Tuple ``(mu, Sigma)`` where mu is shape ``(N,)`` and
        Sigma is a linear operator.
    """
    neg2_eta2 = -2.0 * eta2
    mu = dispatch_solve(neg2_eta2, eta1, solver)
    Sigma = inv(neg2_eta2)
    return mu, Sigma

mean_chol_to_natural(mu: Float[Array, '*batch N'], S_sqrt: Float[Array, '*batch N N']) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Convert mean/variance (Cholesky) to natural parameters.

Given mu and lower-triangular S_sqrt such that Sigma = S_sqrt @ S_sqrt^T:

• eta1 = Sigma^{-1} mu
• eta2 = -0.5 * Sigma^{-1}

Uses the Cholesky factor directly via triangular solves; no solver parameter is exposed because the underlying systems are triangular rather than symmetric/PSD, and iterative strategies (CG, BBMM, PreconditionedCG, MINRES) are not valid here.

Parameters:

Name	Type	Description	Default
mu	Float[Array, '*batch N']	Mean vector, shape (*batch, N).	required
S_sqrt	Float[Array, '*batch N N']	Lower-triangular Cholesky factor, shape (*batch, N, N).	required

Returns:

Type	Description
tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]	Tuple (eta1, eta2) of natural parameters.

Source code in src/gaussx/_expfam/_natural.py

def mean_chol_to_natural(
    mu: Float[Array, "*batch N"],
    S_sqrt: Float[Array, "*batch N N"],
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    r"""Convert mean/variance (Cholesky) to natural parameters.

    Given ``mu`` and lower-triangular ``S_sqrt`` such that
    ``Sigma = S_sqrt @ S_sqrt^T``:

    - ``eta1 = Sigma^{-1} mu``
    - ``eta2 = -0.5 * Sigma^{-1}``

    Uses the Cholesky factor directly via triangular solves; no solver
    parameter is exposed because the underlying systems are triangular
    rather than symmetric/PSD, and iterative strategies (CG, BBMM,
    PreconditionedCG, MINRES) are not valid here.

    Args:
        mu: Mean vector, shape ``(*batch, N)``.
        S_sqrt: Lower-triangular Cholesky factor, shape ``(*batch, N, N)``.

    Returns:
        Tuple ``(eta1, eta2)`` of natural parameters.
    """

    def _core(mu_s: Float[Array, " N"], s_sqrt_s: Float[Array, "N N"]):
        # eta1 = Sigma^{-1} mu = S_sqrt^{-T} S_sqrt^{-1} mu via cho_solve.
        eta1_s = jax.scipy.linalg.cho_solve((s_sqrt_s, True), mu_s)
        # eta2 = -0.5 * Sigma^{-1}, computed by a single matrix cho_solve.
        N = s_sqrt_s.shape[0]
        identity = jnp.eye(N, dtype=s_sqrt_s.dtype)
        Sigma_inv = jax.scipy.linalg.cho_solve((s_sqrt_s, True), identity)
        return eta1_s, -0.5 * Sigma_inv

    *batch, N = mu.shape
    if not batch:
        return _core(mu, S_sqrt)
    mu_flat = mu.reshape(-1, N)
    s_flat = S_sqrt.reshape(-1, N, N)
    eta1_flat, eta2_flat = jax.vmap(_core)(mu_flat, s_flat)
    return eta1_flat.reshape(mu.shape), eta2_flat.reshape(S_sqrt.shape)

natural_to_mean_chol(eta1: Float[Array, '*batch N'], eta2: Float[Array, '*batch N N'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Convert natural parameters to mean/variance (Cholesky).

Given eta1 = Lambda @ mu and eta2 = -0.5 * Lambda:

• Sigma = (-2 * eta2)^{-1}
• mu = Sigma @ eta1
• S_sqrt = cholesky(Sigma)

Parameters:

Name	Type	Description	Default
eta1	Float[Array, '*batch N']	Natural location parameter, shape (*batch, N).	required
eta2	Float[Array, '*batch N N']	Natural quadratic parameter, shape (*batch, N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

Type	Description
Float[Array, '*batch N']	Tuple (mu, S_sqrt) where S_sqrt is the lower-triangular
Float[Array, '*batch N N']	Cholesky factor of the covariance.

Source code in src/gaussx/_expfam/_natural.py

def natural_to_mean_chol(
    eta1: Float[Array, "*batch N"],
    eta2: Float[Array, "*batch N N"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    r"""Convert natural parameters to mean/variance (Cholesky).

    Given ``eta1 = Lambda @ mu`` and ``eta2 = -0.5 * Lambda``:

    - ``Sigma = (-2 * eta2)^{-1}``
    - ``mu = Sigma @ eta1``
    - ``S_sqrt = cholesky(Sigma)``

    Args:
        eta1: Natural location parameter, shape ``(*batch, N)``.
        eta2: Natural quadratic parameter, shape ``(*batch, N, N)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(mu, S_sqrt)`` where ``S_sqrt`` is the lower-triangular
        Cholesky factor of the covariance.
    """

    def _core(e1: Float[Array, " N"], e2: Float[Array, "N N"]):
        Lambda_op = lx.MatrixLinearOperator(-2.0 * e2, lx.positive_semidefinite_tag)
        mu_s = dispatch_solve(Lambda_op, e1, solver)
        Sigma = inv(Lambda_op).as_matrix()
        Sigma_op = lx.MatrixLinearOperator(Sigma, lx.positive_semidefinite_tag)
        return mu_s, cholesky(Sigma_op).as_matrix()

    *batch, N = eta1.shape
    if not batch:
        return _core(eta1, eta2)
    eta1_flat = eta1.reshape(-1, N)
    eta2_flat = eta2.reshape(-1, N, N)
    mu_flat, s_flat = jax.vmap(_core)(eta1_flat, eta2_flat)
    return mu_flat.reshape(eta1.shape), s_flat.reshape(eta2.shape)

mean_chol_to_expectation(mu: Float[Array, '*batch N'], S_sqrt: Float[Array, '*batch N N']) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Convert mean/variance (Cholesky) to expectation parameters.

Given mu and S_sqrt (lower-triangular Cholesky of Sigma):

• m1 = mu
• m2 = mu @ mu^T + Sigma = mu @ mu^T + S_sqrt @ S_sqrt^T

Parameters:

Name	Type	Description	Default
mu	Float[Array, '*batch N']	Mean vector, shape (*batch, N).	required
S_sqrt	Float[Array, '*batch N N']	Lower-triangular Cholesky factor, shape (*batch, N, N).	required

Returns:

Type	Description
tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]	Tuple (m1, m2) of expectation parameters.

Source code in src/gaussx/_expfam/_natural.py

def mean_chol_to_expectation(
    mu: Float[Array, "*batch N"],
    S_sqrt: Float[Array, "*batch N N"],
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    r"""Convert mean/variance (Cholesky) to expectation parameters.

    Given ``mu`` and ``S_sqrt`` (lower-triangular Cholesky of ``Sigma``):

    - ``m1 = mu``
    - ``m2 = mu @ mu^T + Sigma = mu @ mu^T + S_sqrt @ S_sqrt^T``

    Args:
        mu: Mean vector, shape ``(*batch, N)``.
        S_sqrt: Lower-triangular Cholesky factor, shape ``(*batch, N, N)``.

    Returns:
        Tuple ``(m1, m2)`` of expectation parameters.
    """
    m1 = mu
    m2 = mu[..., None] * mu[..., None, :] + S_sqrt @ S_sqrt.mT
    return m1, m2

expectation_to_mean_chol(m1: Float[Array, '*batch N'], m2: Float[Array, '*batch N N']) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Convert expectation parameters to mean/variance (Cholesky).

Given m1 = mu and m2 = mu @ mu^T + Sigma:

• mu = m1
• Sigma = m2 - m1 @ m1^T
• S_sqrt = cholesky(Sigma)

No solver parameter is exposed because the only linear-algebra operation is Cholesky factorization, which is structurally fixed.

Parameters:

Name	Type	Description	Default
m1	Float[Array, '*batch N']	First moment (mean), shape (*batch, N).	required
m2	Float[Array, '*batch N N']	Second moment, shape (*batch, N, N).	required

Returns:

Type	Description
Float[Array, '*batch N']	Tuple (mu, S_sqrt) where S_sqrt is the lower-triangular
Float[Array, '*batch N N']	Cholesky factor of the covariance.

Source code in src/gaussx/_expfam/_natural.py

def expectation_to_mean_chol(
    m1: Float[Array, "*batch N"],
    m2: Float[Array, "*batch N N"],
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    r"""Convert expectation parameters to mean/variance (Cholesky).

    Given ``m1 = mu`` and ``m2 = mu @ mu^T + Sigma``:

    - ``mu = m1``
    - ``Sigma = m2 - m1 @ m1^T``
    - ``S_sqrt = cholesky(Sigma)``

    No solver parameter is exposed because the only linear-algebra
    operation is Cholesky factorization, which is structurally fixed.

    Args:
        m1: First moment (mean), shape ``(*batch, N)``.
        m2: Second moment, shape ``(*batch, N, N)``.

    Returns:
        Tuple ``(mu, S_sqrt)`` where ``S_sqrt`` is the lower-triangular
        Cholesky factor of the covariance.
    """

    def _core(m1_s: Float[Array, " N"], m2_s: Float[Array, "N N"]):
        Sigma = m2_s - m1_s[:, None] * m1_s[None, :]
        Sigma_op = lx.MatrixLinearOperator(Sigma, lx.positive_semidefinite_tag)
        return m1_s, cholesky(Sigma_op).as_matrix()

    *batch, N = m1.shape
    if not batch:
        return _core(m1, m2)
    m1_flat = m1.reshape(-1, N)
    m2_flat = m2.reshape(-1, N, N)
    mu_flat, s_flat = jax.vmap(_core)(m1_flat, m2_flat)
    return mu_flat.reshape(m1.shape), s_flat.reshape(m2.shape)

expectation_to_natural(m1: Float[Array, '*batch N'], m2: Float[Array, '*batch N N'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Convert expectation parameters to natural parameters.

Given m1 = mu and m2 = mu @ mu^T + Sigma:

• Sigma = m2 - m1 @ m1^T
• eta1 = Sigma^{-1} @ m1
• eta2 = -0.5 * Sigma^{-1}

Parameters:

Name	Type	Description	Default
m1	Float[Array, '*batch N']	First moment (mean), shape (*batch, N).	required
m2	Float[Array, '*batch N N']	Second moment, shape (*batch, N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

Type	Description
tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]	Tuple (eta1, eta2) of natural parameters.

Source code in src/gaussx/_expfam/_natural.py

def expectation_to_natural(
    m1: Float[Array, "*batch N"],
    m2: Float[Array, "*batch N N"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    r"""Convert expectation parameters to natural parameters.

    Given ``m1 = mu`` and ``m2 = mu @ mu^T + Sigma``:

    - ``Sigma = m2 - m1 @ m1^T``
    - ``eta1 = Sigma^{-1} @ m1``
    - ``eta2 = -0.5 * Sigma^{-1}``

    Args:
        m1: First moment (mean), shape ``(*batch, N)``.
        m2: Second moment, shape ``(*batch, N, N)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(eta1, eta2)`` of natural parameters.
    """

    def _core(m1_s: Float[Array, " N"], m2_s: Float[Array, "N N"]):
        Sigma = m2_s - m1_s[:, None] * m1_s[None, :]
        Sigma_op = lx.MatrixLinearOperator(Sigma, lx.positive_semidefinite_tag)
        eta1_s = dispatch_solve(Sigma_op, m1_s, solver)
        N = m1_s.shape[0]
        identity = jnp.eye(N, dtype=m1_s.dtype)
        Sigma_inv = solve_columns(Sigma_op, identity, solver=solver)
        return eta1_s, -0.5 * Sigma_inv

    *batch, N = m1.shape
    if not batch:
        return _core(m1, m2)
    m1_flat = m1.reshape(-1, N)
    m2_flat = m2.reshape(-1, N, N)
    eta1_flat, eta2_flat = jax.vmap(_core)(m1_flat, m2_flat)
    return eta1_flat.reshape(m1.shape), eta2_flat.reshape(m2.shape)

natural_to_expectation(eta1: Float[Array, '*batch N'], eta2: Float[Array, '*batch N N'], *, solver: AbstractSolverStrategy | None = None) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Convert natural parameters to expectation parameters.

Given eta1 = Lambda @ mu and eta2 = -0.5 * Lambda:

• Sigma = (-2 * eta2)^{-1}
• mu = Sigma @ eta1
• m1 = mu
• m2 = mu @ mu^T + Sigma

Parameters:

Name	Type	Description	Default
eta1	Float[Array, '*batch N']	Natural location parameter, shape (*batch, N).	required
eta2	Float[Array, '*batch N N']	Natural quadratic parameter, shape (*batch, N, N).	required
solver	AbstractSolverStrategy | None	Optional solver strategy for structured linear algebra. When None, falls back to structural dispatch.	None

Returns:

Type	Description
tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]	Tuple (m1, m2) of expectation parameters.

Source code in src/gaussx/_expfam/_natural.py

def natural_to_expectation(
    eta1: Float[Array, "*batch N"],
    eta2: Float[Array, "*batch N N"],
    *,
    solver: AbstractSolverStrategy | None = None,
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    r"""Convert natural parameters to expectation parameters.

    Given ``eta1 = Lambda @ mu`` and ``eta2 = -0.5 * Lambda``:

    - ``Sigma = (-2 * eta2)^{-1}``
    - ``mu = Sigma @ eta1``
    - ``m1 = mu``
    - ``m2 = mu @ mu^T + Sigma``

    Args:
        eta1: Natural location parameter, shape ``(*batch, N)``.
        eta2: Natural quadratic parameter, shape ``(*batch, N, N)``.
        solver: Optional solver strategy for structured linear algebra.
            When ``None``, falls back to structural dispatch.

    Returns:
        Tuple ``(m1, m2)`` of expectation parameters.
    """

    def _core(e1: Float[Array, " N"], e2: Float[Array, "N N"]):
        Lambda_op = lx.MatrixLinearOperator(-2.0 * e2, lx.positive_semidefinite_tag)
        mu_s = dispatch_solve(Lambda_op, e1, solver)
        Sigma = inv(Lambda_op).as_matrix()
        m2_s = mu_s[:, None] * mu_s[None, :] + Sigma
        return mu_s, m2_s

    *batch, N = eta1.shape
    if not batch:
        return _core(eta1, eta2)
    eta1_flat = eta1.reshape(-1, N)
    eta2_flat = eta2.reshape(-1, N, N)
    m1_flat, m2_flat = jax.vmap(_core)(eta1_flat, eta2_flat)
    return m1_flat.reshape(eta1.shape), m2_flat.reshape(eta2.shape)

log_partition(expfam: GaussianExpFam) -> Float[Array, '']

Log-partition function A(eta).

\[ A(\eta) = -\frac{1}{4} \eta_1^T \eta_2^{-1} \eta_1 - \frac{1}{2} \log|-2\eta_2| + \frac{N}{2} \log 2\pi \]

With the base measure \(h(x) = 1\) (see GaussianExpFam), so that \(\log q(x) = \eta^T T(x) - A(\eta)\) is the normalised log-density.

Parameters:

Name	Type	Description	Default
expfam	GaussianExpFam	Gaussian in natural form.	required

Returns:

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

Source code in src/gaussx/_expfam/_gaussian.py

def log_partition(expfam: GaussianExpFam) -> Float[Array, ""]:
    r"""Log-partition function ``A(eta)``.

    $$
    A(\eta) = -\frac{1}{4} \eta_1^T \eta_2^{-1} \eta_1
              - \frac{1}{2} \log|-2\eta_2|
              + \frac{N}{2} \log 2\pi
    $$

    With the base measure $h(x) = 1$ (see `GaussianExpFam`), so that
    $\log q(x) = \eta^T T(x) - A(\eta)$ is the normalised log-density.

    Args:
        expfam: Gaussian in natural form.

    Returns:
        Scalar log-partition value.
    """
    neg2_eta2 = -2.0 * expfam.eta2
    N = neg2_eta2.in_size()

    # -0.25 * eta1^T @ eta2^{-1} @ eta1
    # eta2^{-1} = (-0.5 Lambda)^{-1} = -2 Sigma
    # So -0.25 * eta1^T @ (-2 Sigma) @ eta1 = 0.5 * eta1^T Sigma eta1
    eta2_inv_eta1 = solve(expfam.eta2, expfam.eta1)
    quad = -0.25 * (expfam.eta1 @ eta2_inv_eta1)

    # -0.5 * log|-2 eta2| = -0.5 * logdet(Lambda)
    ld = -0.5 * logdet(neg2_eta2)

    # Add base measure contribution: N/2 * log(2pi)
    return quad + ld + 0.5 * N * _LOG_2PI

fisher_info(expfam: GaussianExpFam) -> lx.AbstractLinearOperator

Fisher information with respect to the mean, \(F_\mu = \Sigma^{-1}\).

Returns the precision \(\Lambda = -2\eta_2\), which is the Fisher information of the mean parameter \(\mu\) (the \(\mu\) block that natural-gradient descent on a \((\mu, \Sigma)\)-parameterised objective uses). It is not the Fisher information in natural coordinates, \(\nabla^2 A(\eta)\): that is the \((N + N^2)\)-dimensional covariance of \(T(x) = [x, x x^T]\), whose \(\eta_1\)-\(\eta_1\) block is \(\Sigma\), not \(\Lambda\).

Parameters:

Name	Type	Description	Default
expfam	GaussianExpFam	Gaussian in natural form.	required

Returns:

Type	Description
AbstractLinearOperator	Precision operator \(\Lambda = -2\eta_2\).

Source code in src/gaussx/_expfam/_gaussian.py

def fisher_info(
    expfam: GaussianExpFam,
) -> lx.AbstractLinearOperator:
    r"""Fisher information with respect to the mean, $F_\mu = \Sigma^{-1}$.

    Returns the precision $\Lambda = -2\eta_2$, which is the Fisher
    information of the **mean parameter** $\mu$ (the $\mu$ block that
    natural-gradient descent on a $(\mu, \Sigma)$-parameterised objective
    uses). It is **not** the Fisher information in natural coordinates,
    $\nabla^2 A(\eta)$: that is the $(N + N^2)$-dimensional covariance of
    $T(x) = [x, x x^T]$, whose $\eta_1$-$\eta_1$ block is $\Sigma$, not
    $\Lambda$.

    Args:
        expfam: Gaussian in natural form.

    Returns:
        Precision operator $\Lambda = -2\eta_2$.
    """
    # Lambda = -2 * eta2
    return -2.0 * expfam.eta2

sufficient_stats(x: Float[Array, '*batch N']) -> tuple[Float[Array, '*batch N'], Float[Array, '*batch N N']]

Compute sufficient statistics T(x) = [x, x x^T].

Parameters:

Name	Type	Description	Default
x	Float[Array, '*batch N']	Data, shape (*batch, N) with any number of batch axes.	required

Returns:

Type	Description
Float[Array, '*batch N']	Tuple (x, outer_product) where outer_product has shape
Float[Array, '*batch N N']	(*batch, N, N).

Source code in src/gaussx/_expfam/_gaussian.py

def sufficient_stats(
    x: Float[Array, "*batch N"],
) -> tuple[Float[Array, "*batch N"], Float[Array, "*batch N N"]]:
    """Compute sufficient statistics ``T(x) = [x, x x^T]``.

    Args:
        x: Data, shape ``(*batch, N)`` with any number of batch axes.

    Returns:
        Tuple ``(x, outer_product)`` where outer_product has shape
        ``(*batch, N, N)``.
    """
    return x, einsum(x, x, "... i, ... j -> ... i j")

kl_divergence(q: GaussianExpFam, p: GaussianExpFam) -> Float[Array, '']

KL(first ‖ second): KL(q || p) via the Bregman form on natural parameters.

The first argument is the distribution the expectation is taken under, whatever its parameter name: kl_divergence(a, b) = KL(a || b).

Exponential-family expression of the KL divergence in terms of the log-partition A and the natural parameters of q and p. Mathematically equivalent to gaussian_kl.

The current implementation evaluates the Bregman form by routing through to_mean_cov for the natural-gradient term (eta_p - eta_q)^T nabla A(eta_q). The second-moment contraction splits into a quadratic form (operator matvecs) plus gaussx.trace_product, so structured eta2 / Sigma_q operators are never materialized. The benefit relative to gaussian_kl is keeping the gradient flowing in natural-parameter space (suitable inside a natural-gradient loop).

\[ KL(q || p) = A(eta_p) - A(eta_q) - (eta_p - eta_q)^T nabla A(eta_q) \]

Parameters:

Name	Type	Description	Default
q	GaussianExpFam	First Gaussian (the "true" distribution).	required
p	GaussianExpFam	Second Gaussian (the "approximate" distribution).	required

Returns:

Type	Description
Float[Array, '']	Scalar KL divergence.

See Also

gaussian_kl: General KL in mean/covariance form with lineax operators.

Source code in src/gaussx/_expfam/_gaussian.py

def kl_divergence(
    q: GaussianExpFam,
    p: GaussianExpFam,
) -> Float[Array, ""]:
    """KL(first ‖ second): ``KL(q || p)`` via the Bregman form on natural parameters.

    The first argument is the distribution the expectation is taken under,
    whatever its parameter name: ``kl_divergence(a, b) = KL(a || b)``.

    Exponential-family expression of the KL divergence in terms of the
    log-partition ``A`` and the natural parameters of ``q`` and ``p``.
    Mathematically equivalent to
    `gaussian_kl`.

    The current implementation evaluates the Bregman form by routing
    through `to_mean_cov` for the natural-gradient term
    ``(eta_p - eta_q)^T nabla A(eta_q)``. The second-moment contraction
    splits into a quadratic form (operator matvecs) plus
    `gaussx.trace_product`, so structured ``eta2`` / ``Sigma_q``
    operators are never materialized. The benefit relative to
    `gaussian_kl` is keeping the gradient flowing in
    natural-parameter space (suitable inside a natural-gradient loop).

    $$
    KL(q || p) = A(eta_p) - A(eta_q) - (eta_p - eta_q)^T nabla A(eta_q)
    $$

    Args:
        q: First Gaussian (the "true" distribution).
        p: Second Gaussian (the "approximate" distribution).

    Returns:
        Scalar KL divergence.

    See Also:
        `gaussian_kl`: General KL
        in mean/covariance form with lineax operators.
    """
    A_p = log_partition(p)
    A_q = log_partition(q)

    # grad A(eta_q) is the expectation parameters: mu_q w.r.t. eta1 and
    # mu_q mu_q^T + Sigma_q w.r.t. eta2. Build them from (mu_q, Sigma_q).
    # The linear term: (eta_p - eta_q)^T grad A(eta_q)
    # For eta1 part: (eta1_p - eta1_q)^T mu_q
    mu_q, Sigma_q = to_mean_cov(q)

    delta_eta1 = p.eta1 - q.eta1
    linear_eta1 = delta_eta1 @ mu_q

    # For eta2 part: tr((eta2_p - eta2_q) @ (mu mu^T + Sigma))
    # = mu^T (eta2_p - eta2_q) mu + tr(eta2_p Sigma) - tr(eta2_q Sigma).
    # Quadratic form via matvecs + structured trace_product — no
    # materialization of eta2 or Sigma_q.
    quad = mu_q @ (p.eta2.mv(mu_q) - q.eta2.mv(mu_q))
    linear_eta2 = quad + trace_product(p.eta2, Sigma_q) - trace_product(q.eta2, Sigma_q)

    return A_p - A_q - linear_eta1 - linear_eta2

Deprecated aliases

These names still work but warn on use and will be removed in gaussx 0.7.0 (see the deprecation policy).

Structured linear algebra and Gaussian primitives for JAX.

to_natural(mu: Float[Array, ' N'], Sigma: lx.AbstractLinearOperator) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator]

Deprecated alias of gaussx.mean_cov_to_natural.

It takes the mean and covariance, not expectation parameters. It will be removed in gaussx 0.7.0.

Parameters:

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

Returns:

Type	Description
tuple[Float[Array, ' N'], AbstractLinearOperator]	Tuple (eta1, eta2) — natural parameters.

Source code in src/gaussx/_expfam/_gaussian.py

def to_natural(
    mu: Float[Array, " N"],
    Sigma: lx.AbstractLinearOperator,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Deprecated alias of `gaussx.mean_cov_to_natural`.

    It takes the mean and covariance, not expectation parameters. It will
    be removed in gaussx 0.7.0.

    Args:
        mu: Mean vector, shape ``(N,)``.
        Sigma: Covariance operator, shape ``(N, N)``.

    Returns:
        Tuple ``(eta1, eta2)`` — natural parameters.
    """
    warn_deprecated(
        "to_natural is deprecated and will be removed in gaussx 0.7.0: it takes "
        "(mu, Sigma), not expectation "
        "parameters. Use gaussx.mean_cov_to_natural (same result)."
    )
    return mean_cov_to_natural(mu, Sigma)

to_expectation(expfam: GaussianExpFam) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator]

Deprecated alias of to_mean_cov.

Despite its name it returns the mean and covariance (mu, Sigma), not the expectation parameters (mu, mu mu^T + Sigma) (see gaussx.natural_to_expectation). It will be removed in gaussx 0.7.0.

Parameters:

Name	Type	Description	Default
expfam	GaussianExpFam	Gaussian in natural form.	required

Returns:

Type	Description
tuple[Float[Array, ' N'], AbstractLinearOperator]	Tuple (mu, Sigma) — mean vector and covariance operator.

Source code in src/gaussx/_expfam/_gaussian.py

def to_expectation(
    expfam: GaussianExpFam,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Deprecated alias of `to_mean_cov`.

    Despite its name it returns the mean and covariance ``(mu, Sigma)``,
    not the expectation parameters ``(mu, mu mu^T + Sigma)`` (see
    `gaussx.natural_to_expectation`). It will be removed in gaussx 0.7.0.

    Args:
        expfam: Gaussian in natural form.

    Returns:
        Tuple ``(mu, Sigma)`` — mean vector and covariance operator.
    """
    warn_deprecated(
        "to_expectation is deprecated and will be removed in gaussx 0.7.0: it "
        "returns (mu, Sigma), not expectation "
        "parameters. Use gaussx.to_mean_cov (same result), or "
        "gaussx.natural_to_expectation for (mu, mu mu^T + Sigma)."
    )
    return to_mean_cov(expfam)