GaussianDistribution

class gpjax.distributions.GaussianDistribution(loc, scale, validate_args=None)

Bases: Distribution

Multivariate Gaussian distribution for GP predictions.

This is the return type of all predict() methods in GPJax. It wraps a mean vector and a covariance lx.AbstractLinearOperator, providing methods for sampling, computing log-probabilities, and evaluating KL divergences.

The distribution is parameterised as

\[p(\mathbf{x}) = \mathcal{N}(\mathbf{x}; \boldsymbol{\mu}, \mathbf{\Sigma})\]

where \(\boldsymbol{\mu}\) is the loc (mean) vector and \(\mathbf{\Sigma}\) is represented by the scale lx.AbstractLinearOperator.

Parameters:

• loc (Float[Array, " N"]) – Mean vector of the distribution.

• scale (lx.AbstractLinearOperator) – Covariance matrix represented as a Lineax linear operator (e.g. lx.MatrixLinearOperator or lx.DiagonalLinearOperator).

Examples

>>> import jax.numpy as jnp
>>> import lineax as lx
>>> from gpjax.distributions import GaussianDistribution
>>> mu = jnp.array([0.0, 1.0])
>>> cov = lx.MatrixLinearOperator(jnp.eye(2))
>>> dist = GaussianDistribution(loc=mu, scale=cov)
>>> dist.mean
Array([0., 1.], dtype=float32)
>>> dist.variance
Array([1., 1.], dtype=float32)

Expand for references to gpjax.distributions.GaussianDistribution

GaussianDistribution

covariance()

Materialises the full covariance matrix as a dense array.

Returns:

Dense covariance matrix.

Return type:

Float[Array, “N N”]

property covariance_matrix: Float[jaxlib._jax.Array, 'N N'] | Float[ndarray, 'N N']

Property alias for covariance().

entropy()

Calculates the differential entropy of the distribution.

\[H[p] = \tfrac{1}{2}\bigl(N(1 + \ln 2\pi) + \ln|\mathbf{\Sigma}|\bigr)\]

Returns:

Entropy in nats.

Return type:

ScalarFloat

kl_divergence(other)

KL divergence from self to other.

Computes \(\operatorname{KL}[q \| p]\) where self is q and other is p.

Parameters:

other (GaussianDistribution) – The reference distribution p.

Returns:

KL divergence in nats.

Return type:

ScalarFloat

log_prob(y)

Calculates the log pdf of the multivariate Gaussian.

\[\log p(\mathbf{y}) = -\tfrac{1}{2}\bigl[ N\ln 2\pi + \ln|\mathbf{\Sigma}| + (\mathbf{y} - \boldsymbol{\mu})^\top \mathbf{\Sigma}^{-1} (\mathbf{y} - \boldsymbol{\mu}) \bigr]\]

Parameters:

y (Float[Array, " N"]) – Point at which to evaluate the log-density.

Returns:

Log probability.

Return type:

ScalarFloat

property mean: Float[jaxlib._jax.Array, 'N'] | Float[ndarray, 'N']

Calculates the mean.

Expand for references to gpjax.distributions.GaussianDistribution.mean

GaussianDistribution

median()

Calculates the median (equal to the mean for a Gaussian).

Return type:

Float[jaxlib._jax.Array, ‘N’] | Float[ndarray, ‘N’]

mode()

Calculates the mode (equal to the mean for a Gaussian).

Return type:

Float[jaxlib._jax.Array, ‘N’] | Float[ndarray, ‘N’]

sample(key, sample_shape=())

Draw samples from the distribution.

Generates samples via the reparameterisation trick:

\[\mathbf{x} = \boldsymbol{\mu} + \mathbf{L}\mathbf{z}, \quad \mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\]

where \(\mathbf{L}\) is the lower Cholesky factor of the covariance.

Parameters:

• key (KeyArray) – JAX PRNG key.

• sample_shape (tuple of int, optional) – Leading batch dimensions for the samples. Defaults to (), returning a single sample.

Returns:

Array of samples with shape

(*sample_shape, N).

Return type:

Float[Array, “… N”]

stddev()

Calculates the marginal standard deviation.

Return type:

Float[jaxlib._jax.Array, ‘N’] | Float[ndarray, ‘N’]

support = RealVector(Real(), 1)

Parameters:

value (NumLikeT)

Return type:

Array | ndarray | bool | number | bool | int | float | complex

property variance: Float[jaxlib._jax.Array, 'N'] | Float[ndarray, 'N']

Calculates the marginal variance (diagonal of the covariance).

Expand for references to gpjax.distributions.GaussianDistribution.variance

GaussianDistribution