Skip to content

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.

Structured linear algebra and Gaussian primitives for JAX.

MultivariateNormal

Bases: Distribution

Multivariate normal parameterized by a lineax linear operator.

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)
>>> d.log_prob(jnp.ones(3))
Source code in src/gaussx/_distributions/_mvn.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class MultivariateNormal(dist.Distribution):
    """Multivariate normal parameterized by a lineax linear operator.

    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)
        >>> d.log_prob(jnp.ones(3))
    """

    arg_constraints = {"loc": dist.constraints.real_vector}  # noqa: RUF012
    support = dist.constraints.real_vector
    reparametrized_params = ["loc"]  # noqa: RUF012
    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,
        )

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

    @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)

    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."
            )
        L = _cholesky(self.cov_operator)
        shape = sample_shape + self.batch_shape + self.event_shape
        eps = jax.random.normal(key, shape=shape)  # type: ignore[arg-type]
        eps_flat = rearrange(eps, "... D -> (...) D")
        samples_flat = jax.vmap(L.mv)(eps_flat)
        return self.loc + _reshape_samples(samples_flat, shape[:-1])

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

    @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, ""]:
        return gaussian_entropy(self.cov_operator, solver=self.solver)

MultivariateNormalPrecision

Bases: Distribution

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}.

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)
>>> d.log_prob(jnp.ones(3))
Source code in src/gaussx/_distributions/_mvn_prec.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class MultivariateNormalPrecision(dist.Distribution):
    """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}``.

    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)
        >>> d.log_prob(jnp.ones(3))
    """

    arg_constraints = {"loc": dist.constraints.real_vector}  # noqa: RUF012
    support = dist.constraints.real_vector
    reparametrized_params = ["loc"]  # noqa: RUF012
    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,
        )

    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)

    @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)

    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."
            )
        L = _cholesky(self.prec_operator)
        shape = sample_shape + self.batch_shape + self.event_shape
        eps = jax.random.normal(key, shape=shape)  # type: ignore[arg-type]

        def _solve_one(z):
            return _solve(L.T, z)

        eps_flat = rearrange(eps, "... D -> (...) D")
        samples_flat = jax.vmap(_solve_one)(eps_flat)
        return self.loc + _reshape_samples(samples_flat, shape[:-1])

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

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

    def entropy(self) -> Float[Array, ""]:
        n = self.loc.shape[-1]
        ld = self.solver.logdet(self.prec_operator)
        return 0.5 * (n * (1.0 + _LOG_2PI) - ld)

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 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
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))

        x0 = self.m0 + L_P0.mv(jax.random.normal(key_init, (N,)))
        eps_q = jax.random.normal(key_q, (self.n_steps, N))
        eps_r = jax.random.normal(key_r, (self.n_steps, M))

        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 
193
194
195
196
197
198
199
200
201
202
@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 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
@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 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
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

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 
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
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 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
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 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()
        cross = einsum(self.A, covs[:-1], "T i j, T j k -> T i k")
        return pairwise_marginals(means, covs, cross)

    @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 _Q_with_initial(self) -> Float[Array, "T d d"]:
        """``Q`` in the ``Q[0] == P0`` layout used by the UDL helpers."""
        return jnp.concatenate([self.P0[None], self.Q], axis=0)

    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_with_initial())

    @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))
        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()
        residuals = xs[1:] - einsum(self.A, xs[:-1], "T i j, T j -> T i") - self.b
        lp_0 = _chol_log_prob(chol_P0, xs[0] - self.mu0)
        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)
        x0 = self.mu0 + chol_P0 @ jax.random.normal(k0, (d,))
        eps = jax.random.normal(k_rest, (T - 1, d))

        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 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
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 
168
169
170
171
@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 
173
174
175
176
177
@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 
179
180
181
182
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 einsum(self.A, covs[:-1], "T i j, T j k -> T i k")

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 
184
185
186
187
188
189
190
191
192
193
194
195
196
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()
    cross = einsum(self.A, covs[:-1], "T i j, T j k -> T i k")
    return pairwise_marginals(means, covs, cross)

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 
198
199
200
201
202
203
204
205
206
207
208
@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 
218
219
220
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_with_initial())

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 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
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 
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
@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))
    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) -> 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

Returns:

Type Description
Float[Array, '']

Scalar log-probability.
Source code in src/gaussx/_distributions/_gaussian.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def gaussian_log_prob(
    loc: Float[Array, " N"],
    cov_operator: lx.AbstractLinearOperator,
    value: Float[Array, " N"],
    *,
    solver: AbstractSolverStrategy | 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.

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

gaussian_entropy(cov_operator: lx.AbstractLinearOperator, *, solver: AbstractLogdetStrategy | 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

Returns:

Type Description
Float[Array, '']

Scalar entropy.
Source code in src/gaussx/_distributions/_gaussian.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def gaussian_entropy(
    cov_operator: lx.AbstractLinearOperator,
    *,
    solver: AbstractLogdetStrategy | 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.

    Returns:
        Scalar entropy.
    """
    N = cov_operator.in_size()
    ld = dispatch_logdet(cov_operator, solver)
    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 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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) -> Float[Array, '']

KL divergence KL(N(m, S) || N(0, I)).

Special case of dist_kl_divergence 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

Returns:

Type Description
Float[Array, '']

Scalar KL divergence.
See Also

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

Source code in src/gaussx/_distributions/_gaussian.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def kl_standard_normal(
    m: Float[Array, " N"],
    S: lx.AbstractLinearOperator,
    *,
    solver: AbstractLogdetStrategy | None = None,
) -> Float[Array, ""]:
    """KL divergence ``KL(N(m, S) || N(0, I))``.

    Special case of `dist_kl_divergence`
    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.

    Returns:
        Scalar KL divergence.

    See Also:
        `dist_kl_divergence`: 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)
    return 0.5 * (tr_S + mTm - N - ld)

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

KL divergence 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 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def dist_kl_divergence(
    p_loc: Float[Array, " N"],
    p_cov: lx.AbstractLinearOperator,
    q_loc: Float[Array, " N"],
    q_cov: lx.AbstractLinearOperator,
) -> Float[Array, ""]:
    r"""KL divergence ``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

    # 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.
Source code in src/gaussx/_distributions/_conditional.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
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.
    """
    N = loc.shape[0]
    obs_idx = jnp.asarray(obs_idx, dtype=jnp.int32)
    obs_values = jnp.asarray(obs_values, dtype=loc.dtype)

    if obs_idx.ndim != 1:
        raise ValueError("obs_idx must be a 1D array.")
    if obs_values.shape != obs_idx.shape:
        raise ValueError("obs_values must have the same shape as obs_idx.")
    if bool(jnp.any((obs_idx < 0) | (obs_idx >= N))):
        raise ValueError(f"obs_idx must be within bounds [0, {N}).")
    if bool(jnp.any(jnp.diff(jnp.sort(obs_idx)) == 0)):
        raise ValueError("obs_idx must not contain duplicates.")

    # 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 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
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).

 required

Returns:

Type Description
Float[Array, 'B M']

Projection matrix A_X, shape (B, M).
Source code in src/gaussx/_distributions/_project.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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)``.

    Returns:
        Projection matrix A_X, shape ``(B, M)``.
    """
    # 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)

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: A(eta) = -0.25 * eta1^T eta2^{-1} eta1 - 0.5 * log|-2 eta2|
  • Base measure: h(x) = (2 pi)^{-N/2}

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 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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: ``A(eta) = -0.25 * eta1^T eta2^{-1} eta1 - 0.5 * log|-2 eta2|``
    - Base measure: ``h(x) = (2 pi)^{-N/2}``

    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 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@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 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
@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_natural(mu: Float[Array, ' N'], Sigma: lx.AbstractLinearOperator) -> tuple[Float[Array, ' N'], lx.AbstractLinearOperator]

Convert expectation to natural parameters.

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 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def to_natural(
    mu: Float[Array, " N"],
    Sigma: lx.AbstractLinearOperator,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Convert expectation to natural parameters.

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

    Returns:
        Tuple ``(eta1, eta2)`` — natural parameters.
    """
    return mean_cov_to_natural(mu, Sigma)

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

Convert natural to expectation parameters.

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 
77
78
79
80
81
82
83
84
85
86
87
88
def to_expectation(
    expfam: GaussianExpFam,
) -> tuple[Float[Array, " N"], lx.AbstractLinearOperator]:
    """Convert natural to expectation parameters.

    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 meanvar_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 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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
    `meanvar_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_meanvar.

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 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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_meanvar`.

    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

meanvar_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 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def meanvar_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_meanvar(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 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def natural_to_meanvar(
    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)

meanvar_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 
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def meanvar_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_meanvar(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 
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def expectation_to_meanvar(
    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 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
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 
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
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| \]

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 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
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|
    $$

    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 matrix F(eta) = nabla^2 A(eta).

For a Gaussian, the Fisher information in terms of the covariance is Sigma^{-1} (the precision matrix).

Parameters:

Name Type Description Default
expfam GaussianExpFam

Gaussian in natural form.

 required

Returns:

Type Description
AbstractLinearOperator

Precision operator (the Fisher information matrix).
Source code in src/gaussx/_expfam/_gaussian.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def fisher_info(
    expfam: GaussianExpFam,
) -> lx.AbstractLinearOperator:
    r"""Fisher information matrix ``F(eta) = nabla^2 A(eta)``.

    For a Gaussian, the Fisher information in terms of the
    covariance is ``Sigma^{-1}`` (the precision matrix).

    Args:
        expfam: Gaussian in natural form.

    Returns:
        Precision operator (the Fisher information matrix).
    """
    # 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 vector, shape (N,) or batch (B, N).

 required

Returns:

Type Description
Float[Array, '*batch N']

Tuple (x, outer_product) where outer_product has
Float[Array, '*batch N N']

shape (N, N) or (B, N, N).
Source code in src/gaussx/_expfam/_gaussian.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
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 vector, shape ``(N,)`` or batch ``(B, N)``.

    Returns:
        Tuple ``(x, outer_product)`` where outer_product has
        shape ``(N, N)`` or ``(B, N, N)``.
    """
    if x.ndim == 1:
        return x, jnp.outer(x, x)
    # Batched: (B, N) -> (B, N, N)
    return x, einsum(x, x, "b i, b j -> b i j")

kl_divergence(q: GaussianExpFam, p: GaussianExpFam) -> Float[Array, '']

KL divergence KL(q || p) via the Bregman-divergence form on natural parameters.

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 dist_kl_divergence.

The current implementation evaluates the Bregman form by routing through to_expectation 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 dist_kl_divergence 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

dist_kl_divergence: General KL in mean/covariance form with lineax operators.

Source code in src/gaussx/_expfam/_gaussian.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def kl_divergence(
    q: GaussianExpFam,
    p: GaussianExpFam,
) -> Float[Array, ""]:
    """KL divergence ``KL(q || p)`` via the Bregman-divergence form on
    natural parameters.

    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
    `dist_kl_divergence`.

    The current implementation evaluates the Bregman form by routing
    through `to_expectation` 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
    `dist_kl_divergence` 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:
        `dist_kl_divergence`: General KL
        in mean/covariance form with lineax operators.
    """
    A_p = log_partition(p)
    A_q = log_partition(q)

    # grad A(eta_q) w.r.t eta1 = mu_q, w.r.t eta2 = mu_q mu_q^T + 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_expectation(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