gpjax 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

gpjax/kernels/stationary/rational_quadratic.py

@@ -23,7 +23,6 @@ from gpjax.kernels.computations import (
 )
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import squared_distance
-from gpjax.parameters import PositiveReal
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -70,17 +69,15 @@ class RationalQuadratic(StationaryKernel):
             compute_engine: The computation engine that the kernel uses to compute the
                 covariance matrix.
         """
-        if isinstance(alpha, nnx.Variable):
-            self.alpha = alpha
-        else:
-            self.alpha = PositiveReal(alpha)
+        self.alpha = alpha
 
         super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine)
 
     def __call__(self, x: Float[Array, " D"], y: Float[Array, " D"]) -> ScalarFloat:
         x = self.slice_input(x) / self.lengthscale.value
         y = self.slice_input(y) / self.lengthscale.value
-        K = self.variance.value * (
-            1 + 0.5 * squared_distance(x, y) / self.alpha.value
-        ) ** (-self.alpha.value)
+        alpha_val = self.alpha.value if hasattr(self.alpha, "value") else self.alpha
+        K = self.variance.value * (1 + 0.5 * squared_distance(x, y) / alpha_val) ** (
+            -alpha_val
+        )
         return K.squeeze()

gpjax/likelihoods.py

@@ -29,7 +29,6 @@ from gpjax.integrators import (
 )
 from gpjax.parameters import (
     NonNegativeReal,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -59,27 +58,27 @@ class AbstractLikelihood(nnx.Module):
         self.num_datapoints = num_datapoints
         self.integrator = integrator
 
-    def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> npd.Distribution:
+    def __call__(
+        self, dist: tp.Union[npd.MultivariateNormal, GaussianDistribution]
+    ) -> npd.Distribution:
         r"""Evaluate the likelihood function at a given predictive distribution.
 
         Args:
-            *args (Any): Arguments to be passed to the likelihood's `predict` method.
-            **kwargs (Any): Keyword arguments to be passed to the likelihood's
-                `predict` method.
+            dist: The predictive distribution to evaluate the likelihood at.
 
         Returns:
             The predictive distribution.
         """
-        return self.predict(*args, **kwargs)
+        return self.predict(dist)
 
     @abc.abstractmethod
-    def predict(self, *args: tp.Any, **kwargs: tp.Any) -> npd.Distribution:
+    def predict(
+        self, dist: tp.Union[npd.MultivariateNormal, GaussianDistribution]
+    ) -> npd.Distribution:
         r"""Evaluate the likelihood function at a given predictive distribution.
 
         Args:
-            *args (Any): Arguments to be passed to the likelihood's `predict` method.
-            **kwargs (Any): Keyword arguments to be passed to the likelihood's
-                `predict` method.
+            dist: The predictive distribution to evaluate the likelihood at.
 
         Returns:
             npd.Distribution: The predictive distribution.
@@ -133,9 +132,7 @@ class Gaussian(AbstractLikelihood):
     def __init__(
         self,
         num_datapoints: int,
-        obs_stddev: tp.Union[
-            ScalarFloat, Float[Array, "#N"], NonNegativeReal, Static
-        ] = 1.0,
+        obs_stddev: tp.Union[ScalarFloat, Float[Array, "#N"], NonNegativeReal] = 1.0,
         integrator: AbstractIntegrator = AnalyticalGaussianIntegrator(),
     ):
         r"""Initializes the Gaussian likelihood.
@@ -148,7 +145,7 @@ class Gaussian(AbstractLikelihood):
                 likelihoods. Must be an instance of `AbstractIntegrator`. For the Gaussian likelihood, this defaults to
                 the `AnalyticalGaussianIntegrator`, as the expected log likelihood can be computed analytically.
         """
-        if not isinstance(obs_stddev, (NonNegativeReal, Static)):
+        if not isinstance(obs_stddev, NonNegativeReal):
             obs_stddev = NonNegativeReal(jnp.asarray(obs_stddev))
         self.obs_stddev = obs_stddev
 

gpjax/linalg/utils.py

@@ -1,5 +1,8 @@
 """Utility functions for the linear algebra module."""
 
+import jax.numpy as jnp
+from jaxtyping import Array
+
 from gpjax.linalg.operators import LinearOperator
 
 
@@ -31,3 +34,32 @@ def psd(A: LinearOperator) -> LinearOperator:
         A.annotations = set()
     A.annotations.add(PSD)
     return A
+
+
+def add_jitter(matrix: Array, jitter: float | Array = 1e-6) -> Array:
+    """Add jitter to the diagonal of a matrix for numerical stability.
+
+    This function adds a small positive value (jitter) to the diagonal elements
+    of a square matrix to improve numerical stability, particularly for
+    Cholesky decompositions and matrix inversions.
+
+    Args:
+        matrix: A square matrix to which jitter will be added.
+        jitter: The jitter value to add to the diagonal. Defaults to 1e-6.
+
+    Returns:
+        The matrix with jitter added to its diagonal.
+
+    Examples:
+        >>> import jax.numpy as jnp
+        >>> from gpjax.linalg.utils import add_jitter
+        >>> matrix = jnp.array([[1.0, 0.5], [0.5, 1.0]])
+        >>> jittered_matrix = add_jitter(matrix, jitter=0.01)
+    """
+    if matrix.ndim != 2:
+        raise ValueError(f"Expected 2D matrix, got {matrix.ndim}D array")
+
+    if matrix.shape[0] != matrix.shape[1]:
+        raise ValueError(f"Expected square matrix, got shape {matrix.shape}")
+
+    return matrix + jnp.eye(matrix.shape[0]) * jitter

gpjax/mean_functions.py

@@ -27,8 +27,6 @@ from jaxtyping import (
 
 from gpjax.parameters import (
     Parameter,
-    Real,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -132,12 +130,12 @@ class Constant(AbstractMeanFunction):
 
     def __init__(
         self,
-        constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0,
+        constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter] = 0.0,
     ):
-        if isinstance(constant, Parameter) or isinstance(constant, Static):
+        if isinstance(constant, Parameter):
             self.constant = constant
         else:
-            self.constant = Real(jnp.array(constant))
+            self.constant = jnp.array(constant)
 
     def __call__(self, x: Num[Array, "N D"]) -> Float[Array, "N O"]:
         r"""Evaluate the mean function at the given points.
@@ -148,7 +146,10 @@ class Constant(AbstractMeanFunction):
         Returns:
             Float[Array, "1"]: The evaluated mean function.
         """
-        return jnp.ones((x.shape[0], 1)) * self.constant.value
+        if isinstance(self.constant, Parameter):
+            return jnp.ones((x.shape[0], 1)) * self.constant.value
+        else:
+            return jnp.ones((x.shape[0], 1)) * self.constant
 
 
 class Zero(Constant):
@@ -160,7 +161,7 @@ class Zero(Constant):
     """
 
     def __init__(self):
-        super().__init__(constant=Static(jnp.array(0.0)))
+        super().__init__(constant=0.0)
 
 
 class CombinationMeanFunction(AbstractMeanFunction):
@@ -175,7 +176,7 @@ class CombinationMeanFunction(AbstractMeanFunction):
         super().__init__(**kwargs)
 
         # Add means to a list, flattening out instances of this class therein, as in GPFlow kernels.
-        items_list: list[AbstractMeanFunction] = []
+        items_list: list[AbstractMeanFunction] = nnx.List([])
 
         for item in means:
             if not isinstance(item, AbstractMeanFunction):

gpjax/objectives.py

@@ -20,6 +20,7 @@ from gpjax.linalg import (
     psd,
     solve,
 )
+from gpjax.linalg.utils import add_jitter
 from gpjax.typing import (
     Array,
     ScalarFloat,
@@ -97,7 +98,7 @@ def conjugate_mll(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat:
 
     # Σ = (Kxx + Io²) = LLᵀ
     Kxx = posterior.prior.kernel.gram(x)
-    Kxx_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * posterior.prior.jitter
+    Kxx_dense = add_jitter(Kxx.to_dense(), posterior.prior.jitter)
     Sigma_dense = Kxx_dense + jnp.eye(Kxx.shape[0]) * obs_noise
     Sigma = psd(Dense(Sigma_dense))
 
@@ -213,7 +214,7 @@ def log_posterior_density(
 
     # Gram matrix
     Kxx = posterior.prior.kernel.gram(x)
-    Kxx_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * posterior.prior.jitter
+    Kxx_dense = add_jitter(Kxx.to_dense(), posterior.prior.jitter)
     Kxx = psd(Dense(Kxx_dense))
     Lx = lower_cholesky(Kxx)
 
@@ -349,7 +350,7 @@ def collapsed_elbo(variational_family: VF, data: Dataset) -> ScalarFloat:
     noise = variational_family.posterior.likelihood.obs_stddev.value**2
     z = variational_family.inducing_inputs.value
     Kzz = kernel.gram(z)
-    Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * variational_family.jitter
+    Kzz_dense = add_jitter(Kzz.to_dense(), variational_family.jitter)
     Kzz = psd(Dense(Kzz_dense))
     Kzx = kernel.cross_covariance(z, x)
     Kxx_diag = vmap(kernel, in_axes=(0, 0))(x, x)

gpjax/parameters.py

@@ -122,16 +122,6 @@ class SigmoidBounded(Parameter[T]):
             )
 
 
-class Static(nnx.Variable[T]):
-    """Static parameter that is not trainable."""
-
-    def __init__(self, value: T, tag: ParameterTag = "static", **kwargs):
-        _check_is_arraylike(value)
-
-        super().__init__(value=jnp.asarray(value), tag=tag, **kwargs)
-        self._tag = tag
-
-
 class LowerTriangular(Parameter[T]):
     """Parameter that is a lower triangular matrix."""
 

gpjax/variational_families.py

@@ -40,11 +40,11 @@ from gpjax.linalg import (
     psd,
     solve,
 )
+from gpjax.linalg.utils import add_jitter
 from gpjax.mean_functions import AbstractMeanFunction
 from gpjax.parameters import (
     LowerTriangular,
     Real,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -110,11 +110,10 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]):
         inducing_inputs: tp.Union[
             Float[Array, "N D"],
             Real,
-            Static,
         ],
         jitter: ScalarFloat = 1e-6,
     ):
-        if not isinstance(inducing_inputs, (Real, Static)):
+        if not isinstance(inducing_inputs, Real):
             inducing_inputs = Real(inducing_inputs)
 
         self.inducing_inputs = inducing_inputs
@@ -177,25 +176,31 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
                 approximation and the GP prior.
         """
         # Unpack variational parameters
-        mu = self.variational_mean.value
-        sqrt = self.variational_root_covariance.value
-        z = self.inducing_inputs.value
+        variational_mean = self.variational_mean.value
+        variational_sqrt = self.variational_root_covariance.value
+        inducing_inputs = self.inducing_inputs.value
 
         # Unpack mean function and kernel
         mean_function = self.posterior.prior.mean_function
         kernel = self.posterior.prior.kernel
 
-        muz = mean_function(z)
-        Kzz = kernel.gram(z)
-        Kzz = psd(Dense(Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter))
+        inducing_mean = mean_function(inducing_inputs)
+        Kzz = kernel.gram(inducing_inputs)
+        Kzz = psd(Dense(add_jitter(Kzz.to_dense(), self.jitter)))
 
-        sqrt = Triangular(sqrt)
-        S = sqrt @ sqrt.T
+        variational_sqrt_triangular = Triangular(variational_sqrt)
+        variational_covariance = (
+            variational_sqrt_triangular @ variational_sqrt_triangular.T
+        )
 
-        qu = GaussianDistribution(loc=jnp.atleast_1d(mu.squeeze()), scale=S)
-        pu = GaussianDistribution(loc=jnp.atleast_1d(muz.squeeze()), scale=Kzz)
+        q_inducing = GaussianDistribution(
+            loc=jnp.atleast_1d(variational_mean.squeeze()), scale=variational_covariance
+        )
+        p_inducing = GaussianDistribution(
+            loc=jnp.atleast_1d(inducing_mean.squeeze()), scale=Kzz
+        )
 
-        return qu.kl_divergence(pu)
+        return q_inducing.kl_divergence(p_inducing)
 
     def predict(self, test_inputs: Float[Array, "N D"]) -> GaussianDistribution:
         r"""Compute the predictive distribution of the GP at the test inputs t.
@@ -215,26 +220,26 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
                 the test inputs.
         """
         # Unpack variational parameters
-        mu = self.variational_mean.value
-        sqrt = self.variational_root_covariance.value
-        z = self.inducing_inputs.value
+        variational_mean = self.variational_mean.value
+        variational_sqrt = self.variational_root_covariance.value
+        inducing_inputs = self.inducing_inputs.value
 
         # Unpack mean function and kernel
         mean_function = self.posterior.prior.mean_function
         kernel = self.posterior.prior.kernel
 
-        Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = kernel.gram(inducing_inputs)
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
-        muz = mean_function(z)
+        inducing_mean = mean_function(inducing_inputs)
 
         # Unpack test inputs
-        t = test_inputs
+        test_points = test_inputs
 
-        Ktt = kernel.gram(t)
-        Kzt = kernel.cross_covariance(z, t)
-        mut = mean_function(t)
+        Ktt = kernel.gram(test_points)
+        Kzt = kernel.cross_covariance(inducing_inputs, test_points)
+        test_mean = mean_function(test_points)
 
         # Lz⁻¹ Kzt
         Lz_inv_Kzt = solve(Lz, Kzt)
@@ -243,10 +248,10 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt)
 
         # Ktz Kzz⁻¹ sqrt
-        Ktz_Kzz_inv_sqrt = jnp.matmul(Kzz_inv_Kzt.T, sqrt)
+        Ktz_Kzz_inv_sqrt = jnp.matmul(Kzz_inv_Kzt.T, variational_sqrt)
 
         # μt + Ktz Kzz⁻¹ (μ - μz)
-        mean = mut + jnp.matmul(Kzz_inv_Kzt.T, mu - muz)
+        mean = test_mean + jnp.matmul(Kzz_inv_Kzt.T, variational_mean - inducing_mean)
 
         # Ktt - Ktz Kzz⁻¹ Kzt  +  Ktz Kzz⁻¹ S Kzz⁻¹ Kzt  [recall S = sqrt sqrtᵀ]
         covariance = (
@@ -254,7 +259,10 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(Ktz_Kzz_inv_sqrt, Ktz_Kzz_inv_sqrt.T)
         )
-        covariance += jnp.eye(covariance.shape[0]) * self.jitter
+        if hasattr(covariance, "to_dense"):
+            covariance = covariance.to_dense()
+        covariance = add_jitter(covariance, self.jitter)
+        covariance = Dense(covariance)
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -329,7 +337,7 @@ class WhitenedVariationalGaussian(VariationalGaussian[L]):
         kernel = self.posterior.prior.kernel
 
         Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
 
@@ -355,7 +363,10 @@ class WhitenedVariationalGaussian(VariationalGaussian[L]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(Ktz_Lz_invT_sqrt, Ktz_Lz_invT_sqrt.T)
         )
-        covariance += jnp.eye(covariance.shape[0]) * self.jitter
+        if hasattr(covariance, "to_dense"):
+            covariance = covariance.to_dense()
+        covariance = add_jitter(covariance, self.jitter)
+        covariance = Dense(covariance)
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -390,8 +401,8 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
         if natural_matrix is None:
             natural_matrix = -0.5 * jnp.eye(self.num_inducing)
 
-        self.natural_vector = Static(natural_vector)
-        self.natural_matrix = Static(natural_matrix)
+        self.natural_vector = Real(natural_vector)
+        self.natural_matrix = Real(natural_matrix)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the KL-divergence between our current variational approximation
@@ -422,7 +433,7 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
 
         # S⁻¹ = -2θ₂
         S_inv = -2 * natural_matrix
-        S_inv += jnp.eye(m) * self.jitter
+        S_inv = add_jitter(S_inv, self.jitter)
 
         # Compute L⁻¹, where LLᵀ = S, via a trick found in the NumPyro source code and https://nbviewer.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril:
         sqrt_inv = jnp.swapaxes(
@@ -441,7 +452,7 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
 
         qu = GaussianDistribution(loc=jnp.atleast_1d(mu.squeeze()), scale=S)
@@ -476,7 +487,7 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
 
         # S⁻¹ = -2θ₂
         S_inv = -2 * natural_matrix
-        S_inv += jnp.eye(m) * self.jitter
+        S_inv = add_jitter(S_inv, self.jitter)
 
         # Compute L⁻¹, where LLᵀ = S, via a trick found in the NumPyro source code and https://nbviewer.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril:
         sqrt_inv = jnp.swapaxes(
@@ -493,7 +504,7 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
         mu = jnp.matmul(S, natural_vector)
 
         Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
@@ -520,7 +531,10 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(Ktz_Kzz_inv_L, Ktz_Kzz_inv_L.T)
         )
-        covariance += jnp.eye(covariance.shape[0]) * self.jitter
+        if hasattr(covariance, "to_dense"):
+            covariance = covariance.to_dense()
+        covariance = add_jitter(covariance, self.jitter)
+        covariance = Dense(covariance)
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -556,8 +570,8 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
         if expectation_matrix is None:
             expectation_matrix = jnp.eye(self.num_inducing)
 
-        self.expectation_vector = Static(expectation_vector)
-        self.expectation_matrix = Static(expectation_matrix)
+        self.expectation_vector = Real(expectation_vector)
+        self.expectation_matrix = Real(expectation_matrix)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Evaluate the prior KL-divergence.
@@ -595,12 +609,12 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
         # S = η₂ - η₁ η₁ᵀ
         S = expectation_matrix - jnp.outer(mu, mu)
         S = psd(Dense(S))
-        S_dense = S.to_dense() + jnp.eye(S.shape[0]) * self.jitter
+        S_dense = add_jitter(S.to_dense(), self.jitter)
         S = psd(Dense(S_dense))
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
 
         qu = GaussianDistribution(loc=jnp.atleast_1d(mu.squeeze()), scale=S)
@@ -640,14 +654,14 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
 
         # S = η₂ - η₁ η₁ᵀ
         S = expectation_matrix - jnp.matmul(mu, mu.T)
-        S = Dense(S + jnp.eye(S.shape[0]) * self.jitter)
+        S = Dense(add_jitter(S, self.jitter))
         S = psd(S)
 
         # S = sqrt sqrtᵀ
         sqrt = lower_cholesky(S)
 
         Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
@@ -677,7 +691,10 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(Ktz_Kzz_inv_sqrt, Ktz_Kzz_inv_sqrt.T)
         )
-        covariance += jnp.eye(covariance.shape[0]) * self.jitter
+        if hasattr(covariance, "to_dense"):
+            covariance = covariance.to_dense()
+        covariance = add_jitter(covariance, self.jitter)
+        covariance = Dense(covariance)
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -734,7 +751,7 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]):
 
         Kzx = kernel.cross_covariance(z, x)
         Kzz = kernel.gram(z)
-        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
         Kzz = psd(Dense(Kzz_dense))
 
         # Lz Lzᵀ = Kzz
@@ -780,7 +797,10 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]):
             - jnp.matmul(Lz_inv_Kzt.T, Lz_inv_Kzt)
             + jnp.matmul(L_inv_Lz_inv_Kzt.T, L_inv_Lz_inv_Kzt)
         )
-        covariance += jnp.eye(covariance.shape[0]) * self.jitter
+        if hasattr(covariance, "to_dense"):
+            covariance = covariance.to_dense()
+        covariance = add_jitter(covariance, self.jitter)
+        covariance = Dense(covariance)
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance

{gpjax-0.12.0.dist-info → gpjax-0.13.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.12.0
+Version: 0.13.0
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -11,15 +11,14 @@ License-File: LICENSE.txt
 Keywords: gaussian-processes jax machine-learning bayesian
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
-Requires-Python: <=3.13,>=3.10
+Requires-Python: >=3.11
 Requires-Dist: beartype>0.16.1
-Requires-Dist: flax>=0.10.0
+Requires-Dist: flax>=0.12.0
 Requires-Dist: jax>=0.5.0
 Requires-Dist: jaxlib>=0.5.0
 Requires-Dist: jaxtyping>0.2.10
@@ -60,7 +59,7 @@ Requires-Dist: mkdocs-jupyter>=0.24.3; extra == 'docs'
 Requires-Dist: mkdocs-literate-nav>=0.6.0; extra == 'docs'
 Requires-Dist: mkdocs-material>=9.5.12; extra == 'docs'
 Requires-Dist: mkdocs>=1.5.3; extra == 'docs'
-Requires-Dist: mkdocstrings[python]<0.28.0; extra == 'docs'
+Requires-Dist: mkdocstrings[python]<0.31.0; extra == 'docs'
 Requires-Dist: nbconvert>=7.16.2; extra == 'docs'
 Requires-Dist: networkx>=3.0; extra == 'docs'
 Requires-Dist: pandas>=1.5.3; extra == 'docs'
@@ -80,6 +79,7 @@ Description-Content-Type: text/markdown
 [![CodeFactor](https://www.codefactor.io/repository/github/jaxgaussianprocesses/gpjax/badge)](https://www.codefactor.io/repository/github/jaxgaussianprocesses/gpjax)
 [![Netlify Status](https://api.netlify.com/api/v1/badges/d3950e6f-321f-4508-9e52-426b5dae2715/deploy-status)](https://app.netlify.com/sites/endearing-crepe-c2d5fe/deploys)
 [![PyPI version](https://badge.fury.io/py/GPJax.svg)](https://badge.fury.io/py/GPJax)
+[![Conda Version](https://img.shields.io/conda/vn/conda-forge/gpjax.svg)](https://anaconda.org/conda-forge/gpjax)
 [![DOI](https://joss.theoj.org/papers/10.21105/joss.04455/status.svg)](https://doi.org/10.21105/joss.04455)
 [![Downloads](https://pepy.tech/badge/gpjax)](https://pepy.tech/project/gpjax)
 [![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
@@ -126,18 +126,9 @@ Channel](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw
 where we can discuss the development of GPJax and broader support for Gaussian
 process modelling.
 
-
-## Governance
-
-GPJax was founded by [Thomas Pinder](https://github.com/thomaspinder). Today, the
-project's gardeners are [daniel-dodd@](https://github.com/daniel-dodd),
-[henrymoss@](https://github.com/henrymoss), [st--@](https://github.com/st--), and
-[thomaspinder@](https://github.com/thomaspinder), listed in alphabetical order. The full
-governance structure of GPJax is detailed [here](docs/GOVERNANCE.md). We appreciate all
-[the contributors to
-GPJax](https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors) who have
-helped to shape GPJax into the package it is today.
-
+We appreciate all [the contributors to
+GPJax](https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors) who have helped to shape
+GPJax into the package it is today.
 
 # Supported methods and interfaces
 
@@ -183,13 +174,21 @@ jupytext --to py:percent example.ipynb
 
 ## Stable version
 
-The latest stable version of GPJax can be installed via
-pip:
+The latest stable version of GPJax can be installed from [PyPI](https://pypi.org/project/gpjax/):
 
 ```bash
 pip install gpjax
 ```
 
+or from [conda-forge](https://github.com/conda-forge/gpjax-feedstock):
+
+```bash
+# with Pixi
+pixi add gpjax
+# or with conda
+conda install --channel conda-forge gpjax
+```
+
 > **Note**
 >
 > We recommend you check your installation version:
@@ -208,7 +207,7 @@ pip install gpjax
 >
 > We advise you create virtual environment before installing:
 > ```
-> conda create -n gpjax_experimental python=3.10.0
+> conda create -n gpjax_experimental python=3.11.0
 > conda activate gpjax_experimental
 >  ```
 
@@ -218,13 +217,14 @@ configuration in development mode.
 ```bash
 git clone https://github.com/JaxGaussianProcesses/GPJax.git
 cd GPJax
+uv venv
 uv sync --extra dev
 ```
 
 > We recommend you check your installation passes the supplied unit tests:
 >
 > ```python
-> uv run pytest --beartype-packages='gpjax'
+> uv run poe all-tests
 > ```
 
 # Citing GPJax