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

gpjax/__init__.py

@@ -40,10 +40,9 @@ __license__ = "MIT"
 __description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.12.0"
+__version__ = "0.12.2"
 
 __all__ = [
-    "base",
     "gps",
     "integrators",
     "kernels",
@@ -55,8 +54,6 @@ __all__ = [
     "Dataset",
     "cite",
     "fit",
-    "Module",
-    "param_field",
     "fit_lbfgs",
     "fit_scipy",
 ]

gpjax/fit.py

@@ -48,6 +48,7 @@ def fit(  # noqa: PLR0913
     train_data: Dataset,
     optim: ox.GradientTransformation,
     params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
+    trainable: nnx.filterlib.Filter = Parameter,
     key: KeyArray = jr.PRNGKey(42),
     num_iters: int = 100,
     batch_size: int = -1,
@@ -65,7 +66,7 @@ def fit(  # noqa: PLR0913
         >>> import jax.random as jr
         >>> import optax as ox
         >>> import gpjax as gpx
-        >>> from gpjax.parameters import PositiveReal, Static
+        >>> from gpjax.parameters import PositiveReal
         >>>
         >>> # (1) Create a dataset:
         >>> X = jnp.linspace(0.0, 10.0, 100)[:, None]
@@ -75,10 +76,10 @@ def fit(  # noqa: PLR0913
         >>> class LinearModel(nnx.Module):
         >>>     def __init__(self, weight: float, bias: float):
         >>>         self.weight = PositiveReal(weight)
-        >>>         self.bias = Static(bias)
+        >>>         self.bias = bias
         >>>
         >>>     def __call__(self, x):
-        >>>         return self.weight.value * x + self.bias.value
+        >>>         return self.weight.value * x + self.bias
         >>>
         >>> model = LinearModel(weight=1.0, bias=1.0)
         >>>
@@ -100,6 +101,8 @@ def fit(  # noqa: PLR0913
         train_data (Dataset): The training data to be used for the optimisation.
         optim (GradientTransformation): The Optax optimiser that is to be used for
             learning a parameter set.
+        trainable (nnx.filterlib.Filter): Filter to determine which parameters are trainable.
+            Defaults to nnx.Param (all Parameter instances).
         num_iters (int): The number of optimisation steps to run. Defaults
             to 100.
         batch_size (int): The size of the mini-batch to use. Defaults to -1
@@ -127,7 +130,7 @@ def fit(  # noqa: PLR0913
         _check_verbose(verbose)
 
     # Model state filtering
-    graphdef, params, *static_state = nnx.split(model, Parameter, ...)
+    graphdef, params, *static_state = nnx.split(model, trainable, ...)
 
     # Parameters bijection to unconstrained space
     if params_bijection is not None:
@@ -182,6 +185,7 @@ def fit_scipy(  # noqa: PLR0913
     model: Model,
     objective: Objective,
     train_data: Dataset,
+    trainable: nnx.filterlib.Filter = Parameter,
     max_iters: int = 500,
     verbose: bool = True,
     safe: bool = True,
@@ -210,7 +214,7 @@ def fit_scipy(  # noqa: PLR0913
         _check_verbose(verbose)
 
     # Model state filtering
-    graphdef, params, *static_state = nnx.split(model, Parameter, ...)
+    graphdef, params, *static_state = nnx.split(model, trainable, ...)
 
     # Parameters bijection to unconstrained space
     params = transform(params, DEFAULT_BIJECTION, inverse=True)
@@ -258,6 +262,7 @@ def fit_lbfgs(
     objective: Objective,
     train_data: Dataset,
     params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
+    trainable: nnx.filterlib.Filter = Parameter,
     max_iters: int = 100,
     safe: bool = True,
     max_linesearch_steps: int = 32,
@@ -290,7 +295,7 @@ def fit_lbfgs(
         _check_num_iters(max_iters)
 
     # Model state filtering
-    graphdef, params, *static_state = nnx.split(model, Parameter, ...)
+    graphdef, params, *static_state = nnx.split(model, trainable, ...)
 
     # Parameters bijection to unconstrained space
     if params_bijection is not None:

gpjax/gps.py

@@ -16,9 +16,9 @@
 from abc import abstractmethod
 
 import beartype.typing as tp
-from flax import nnx
 import jax.numpy as jnp
 import jax.random as jr
+from flax import nnx
 from jaxtyping import (
     Float,
     Num,
@@ -35,16 +35,15 @@ from gpjax.likelihoods import (
 )
 from gpjax.linalg import (
     Dense,
-    Identity,
     psd,
     solve,
 )
 from gpjax.linalg.operations import lower_cholesky
+from gpjax.linalg.utils import add_jitter
 from gpjax.mean_functions import AbstractMeanFunction
 from gpjax.parameters import (
     Parameter,
     Real,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -78,7 +77,7 @@ class AbstractPrior(nnx.Module, tp.Generic[M, K]):
         self.mean_function = mean_function
         self.jitter = jitter
 
-    def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> GaussianDistribution:
+    def __call__(self, test_inputs: Num[Array, "N D"]) -> GaussianDistribution:
         r"""Evaluate the Gaussian process at the given points.
 
         The output of this function is a
@@ -91,17 +90,16 @@ class AbstractPrior(nnx.Module, tp.Generic[M, K]):
         `__call__` method and should instead define a `predict` method.
 
         Args:
-            *args (Any): The arguments to pass to the GP's `predict` method.
-            **kwargs (Any): The keyword arguments to pass to the GP's `predict` method.
+            test_inputs: Input locations where the GP should be evaluated.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
                 of the Gaussian process.
         """
-        return self.predict(*args, **kwargs)
+        return self.predict(test_inputs)
 
     @abstractmethod
-    def predict(self, *args: tp.Any, **kwargs: tp.Any) -> GaussianDistribution:
+    def predict(self, test_inputs: Num[Array, "N D"]) -> GaussianDistribution:
         r"""Evaluate the predictive distribution.
 
         Compute the latent function's multivariate normal distribution for a
@@ -109,8 +107,7 @@ class AbstractPrior(nnx.Module, tp.Generic[M, K]):
         this method must be implemented.
 
         Args:
-            *args (Any): Arguments to the predict method.
-            **kwargs (Any): Keyword arguments to the predict method.
+            test_inputs: Input locations where the GP should be evaluated.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
@@ -249,13 +246,12 @@ class Prior(AbstractPrior[M, K]):
             GaussianDistribution: A multivariate normal random variable representation
                 of the Gaussian process.
         """
-        x = test_inputs
-        mx = self.mean_function(x)
-        Kxx = self.kernel.gram(x)
-        Kxx_dense = Kxx.to_dense() + Identity(Kxx.shape).to_dense() * self.jitter
+        mean_at_test = self.mean_function(test_inputs)
+        Kxx = self.kernel.gram(test_inputs)
+        Kxx_dense = add_jitter(Kxx.to_dense(), self.jitter)
         Kxx = psd(Dense(Kxx_dense))
 
-        return GaussianDistribution(jnp.atleast_1d(mx.squeeze()), Kxx)
+        return GaussianDistribution(jnp.atleast_1d(mean_at_test.squeeze()), Kxx)
 
     def sample_approx(
         self,
@@ -359,7 +355,9 @@ class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
         self.likelihood = likelihood
         self.jitter = jitter
 
-    def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> GaussianDistribution:
+    def __call__(
+        self, test_inputs: Num[Array, "N D"], train_data: Dataset
+    ) -> GaussianDistribution:
         r"""Evaluate the Gaussian process posterior at the given points.
 
         The output of this function is a
@@ -368,28 +366,30 @@ class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
         evaluated and the distribution can be sampled.
 
         Under the hood, `__call__` is calling the objects `predict` method. For this
-        reasons, classes inheriting the `AbstractPrior` class, should not overwrite the
+        reasons, classes inheriting the `AbstractPosterior` class, should not overwrite the
         `__call__` method and should instead define a `predict` method.
 
         Args:
-            *args (Any): The arguments to pass to the GP's `predict` method.
-            **kwargs (Any): The keyword arguments to pass to the GP's `predict` method.
+            test_inputs: Input locations where the GP should be evaluated.
+            train_data: Training dataset to condition on.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
                 of the Gaussian process.
         """
-        return self.predict(*args, **kwargs)
+        return self.predict(test_inputs, train_data)
 
     @abstractmethod
-    def predict(self, *args: tp.Any, **kwargs: tp.Any) -> GaussianDistribution:
+    def predict(
+        self, test_inputs: Num[Array, "N D"], train_data: Dataset
+    ) -> GaussianDistribution:
         r"""Compute the latent function's multivariate normal distribution for a
-        given set of parameters. For any class inheriting the `AbstractPrior` class,
+        given set of parameters. For any class inheriting the `AbstractPosterior` class,
         this method must be implemented.
 
         Args:
-            *args (Any): Arguments to the predict method.
-            **kwargs (Any): Keyword arguments to the predict method.
+            test_inputs: Input locations where the GP should be evaluated.
+            train_data: Training dataset to condition on.
 
         Returns:
             GaussianDistribution: A multivariate normal random variable representation
@@ -503,22 +503,24 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
 
         # Precompute Gram matrix, Kxx, at training inputs, x
         Kxx = self.prior.kernel.gram(x)
-        Kxx_dense = Kxx.to_dense() + Identity(Kxx.shape).to_dense() * self.jitter
+        Kxx_dense = add_jitter(Kxx.to_dense(), self.jitter)
         Kxx = Dense(Kxx_dense)
 
         Sigma_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * obs_noise
         Sigma = psd(Dense(Sigma_dense))
+        L_sigma = lower_cholesky(Sigma)
 
         mean_t = self.prior.mean_function(t)
         Ktt = self.prior.kernel.gram(t)
         Kxt = self.prior.kernel.cross_covariance(x, t)
-        Sigma_inv_Kxt = solve(Sigma, Kxt)
 
-        mean = mean_t + jnp.matmul(Sigma_inv_Kxt.T, y - mx)
+        L_inv_Kxt = solve(L_sigma, Kxt)
+        L_inv_y_diff = solve(L_sigma, y - mx)
+
+        mean = mean_t + jnp.matmul(L_inv_Kxt.T, L_inv_y_diff)
 
-        # Ktt  -  Ktx (Kxx + Io²)⁻¹ Kxt, TODO: Take advantage of covariance structure to compute Schur complement more efficiently.
-        covariance = Ktt.to_dense() - jnp.matmul(Kxt.T, Sigma_inv_Kxt)
-        covariance += jnp.eye(covariance.shape[0]) * self.prior.jitter
+        covariance = Ktt.to_dense() - jnp.matmul(L_inv_Kxt.T, L_inv_Kxt)
+        covariance = add_jitter(covariance, self.prior.jitter)
         covariance = psd(Dense(covariance))
 
         return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), covariance)
@@ -577,7 +579,7 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
 
         obs_var = self.likelihood.obs_stddev.value**2
         Kxx = self.prior.kernel.gram(train_data.X)
-        Sigma = Kxx + jnp.eye(Kxx.shape[0]) * (obs_var + self.jitter)
+        Sigma = Dense(add_jitter(Kxx.to_dense(), obs_var + self.jitter))
         eps = jnp.sqrt(obs_var) * jr.normal(key, [train_data.n, num_samples])
         y = train_data.y - self.prior.mean_function(train_data.X)
         Phi = fourier_feature_fn(train_data.X)
@@ -643,7 +645,7 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
 
         # TODO: static or intermediate?
         self.latent = latent if isinstance(latent, Parameter) else Real(latent)
-        self.key = Static(key)
+        self.key = key
 
     def predict(
         self, test_inputs: Num[Array, "N D"], train_data: Dataset
@@ -675,7 +677,7 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
 
         # Precompute lower triangular of Gram matrix, Lx, at training inputs, x
         Kxx = kernel.gram(x)
-        Kxx_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * self.prior.jitter
+        Kxx_dense = add_jitter(Kxx.to_dense(), self.prior.jitter)
         Kxx = psd(Dense(Kxx_dense))
         Lx = lower_cholesky(Kxx)
 
@@ -698,7 +700,7 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
 
         # Ktt - Ktx Kxx⁻¹ Kxt, TODO: Take advantage of covariance structure to compute Schur complement more efficiently.
         covariance = Ktt.to_dense() - jnp.matmul(Lx_inv_Kxt.T, Lx_inv_Kxt)
-        covariance += jnp.eye(covariance.shape[0]) * self.prior.jitter
+        covariance = add_jitter(covariance, self.prior.jitter)
         covariance = psd(Dense(covariance))
 
         return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), covariance)

gpjax/kernels/approximations/rff.py

@@ -7,7 +7,6 @@ from jaxtyping import Float
 from gpjax.kernels.base import AbstractKernel
 from gpjax.kernels.computations import BasisFunctionComputation
 from gpjax.kernels.stationary.base import StationaryKernel
-from gpjax.parameters import Static
 from gpjax.typing import (
     Array,
     KeyArray,
@@ -66,10 +65,8 @@ class RFF(AbstractKernel):
                     "Please specify the n_dims argument for the base kernel."
                 )
 
-            self.frequencies = Static(
-                self.base_kernel.spectral_density.sample(
-                    key=key, sample_shape=(self.num_basis_fns, n_dims)
-                )
+            self.frequencies = self.base_kernel.spectral_density.sample(
+                key=key, sample_shape=(self.num_basis_fns, n_dims)
             )
         self.name = f"{self.base_kernel.name} (RFF)"
 

gpjax/kernels/base.py

@@ -32,7 +32,6 @@ from gpjax.linalg import LinearOperator
 from gpjax.parameters import (
     Parameter,
     Real,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -221,9 +220,7 @@ class Constant(AbstractKernel):
     def __init__(
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
-        constant: tp.Union[
-            ScalarFloat, Parameter[ScalarFloat], Static[ScalarFloat]
-        ] = jnp.array(0.0),
+        constant: tp.Union[ScalarFloat, Parameter[ScalarFloat]] = jnp.array(0.0),
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),
     ):
         if isinstance(constant, Parameter):

gpjax/kernels/computations/basis_functions.py

@@ -57,7 +57,7 @@ class BasisFunctionComputation(AbstractKernelComputation):
         Returns:
             A matrix of shape $N \times L$ representing the random fourier features where $L = 2M$.
         """
-        frequencies = kernel.frequencies.value
+        frequencies = kernel.frequencies
         scaling_factor = kernel.base_kernel.lengthscale.value
         z = jnp.matmul(x, (frequencies / scaling_factor).T)
         z = jnp.concatenate([jnp.cos(z), jnp.sin(z)], axis=-1)

gpjax/kernels/computations/eigen.py

@@ -42,7 +42,7 @@ class EigenKernelComputation(AbstractKernelComputation):
         # Transform the eigenvalues of the graph Laplacian according to the
         # RBF kernel's SPDE form.
         S = jnp.power(
-            kernel.eigenvalues.value
+            kernel.eigenvalues
             + 2
             * kernel.smoothness.value
             / kernel.lengthscale.value

gpjax/kernels/non_euclidean/graph.py

@@ -30,7 +30,6 @@ from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.parameters import (
     Parameter,
     PositiveReal,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -55,9 +54,9 @@ class GraphKernel(StationaryKernel):
     """
 
     num_vertex: tp.Union[ScalarInt, None]
-    laplacian: Static[Float[Array, "N N"]]
-    eigenvalues: Static[Float[Array, "N 1"]]
-    eigenvectors: Static[Float[Array, "N N"]]
+    laplacian: Float[Array, "N N"]
+    eigenvalues: Float[Array, "N 1"]
+    eigenvectors: Float[Array, "N N"]
     name: str = "Graph Matérn"
 
     def __init__(
@@ -91,11 +90,11 @@ class GraphKernel(StationaryKernel):
         else:
             self.smoothness = PositiveReal(smoothness)
 
-        self.laplacian = Static(laplacian)
-        evals, eigenvectors = jnp.linalg.eigh(self.laplacian.value)
-        self.eigenvectors = Static(eigenvectors)
-        self.eigenvalues = Static(evals.reshape(-1, 1))
-        self.num_vertex = self.eigenvalues.value.shape[0]
+        self.laplacian = laplacian
+        evals, eigenvectors = jnp.linalg.eigh(self.laplacian)
+        self.eigenvectors = eigenvectors
+        self.eigenvalues = evals.reshape(-1, 1)
+        self.num_vertex = self.eigenvalues.shape[0]
 
         super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine)
 
@@ -107,7 +106,7 @@ class GraphKernel(StationaryKernel):
         S,
         **kwargs,
     ):
-        Kxx = (jax_gather_nd(self.eigenvectors.value, x) * S.squeeze()) @ jnp.transpose(
-            jax_gather_nd(self.eigenvectors.value, y)
+        Kxx = (jax_gather_nd(self.eigenvectors, x) * S.squeeze()) @ jnp.transpose(
+            jax_gather_nd(self.eigenvectors, y)
         )  # shape (n,n)
         return Kxx.squeeze()

gpjax/kernels/nonstationary/arccosine.py

@@ -25,7 +25,6 @@ from gpjax.kernels.computations import (
 )
 from gpjax.parameters import (
     NonNegativeReal,
-    PositiveReal,
 )
 from gpjax.typing import (
     Array,
@@ -82,30 +81,13 @@ class ArcCosine(AbstractKernel):
 
         self.order = order
 
-        if isinstance(weight_variance, nnx.Variable):
-            self.weight_variance = weight_variance
-        else:
-            self.weight_variance = PositiveReal(weight_variance)
-            if tp.TYPE_CHECKING:
-                self.weight_variance = tp.cast(
-                    PositiveReal[WeightVariance], self.weight_variance
-                )
+        self.weight_variance = weight_variance
 
         if isinstance(variance, nnx.Variable):
             self.variance = variance
         else:
             self.variance = NonNegativeReal(variance)
-            if tp.TYPE_CHECKING:
-                self.variance = tp.cast(NonNegativeReal[ScalarArray], self.variance)
-
-        if isinstance(bias_variance, nnx.Variable):
-            self.bias_variance = bias_variance
-        else:
-            self.bias_variance = PositiveReal(bias_variance)
-            if tp.TYPE_CHECKING:
-                self.bias_variance = tp.cast(
-                    PositiveReal[ScalarArray], self.bias_variance
-                )
+        self.bias_variance = bias_variance
 
         self.name = f"ArcCosine (order {self.order})"
 
@@ -141,7 +123,17 @@ class ArcCosine(AbstractKernel):
         Returns:
             ScalarFloat: The value of the weighted product between the two arguments``.
         """
-        return jnp.inner(self.weight_variance.value * x, y) + self.bias_variance.value
+        weight_var = (
+            self.weight_variance.value
+            if hasattr(self.weight_variance, "value")
+            else self.weight_variance
+        )
+        bias_var = (
+            self.bias_variance.value
+            if hasattr(self.bias_variance, "value")
+            else self.bias_variance
+        )
+        return jnp.inner(weight_var * x, y) + bias_var
 
     def _J(self, theta: ScalarFloat) -> ScalarFloat:
         r"""Evaluate the angular dependency function corresponding to the desired order.

gpjax/kernels/nonstationary/polynomial.py

@@ -69,12 +69,9 @@ class Polynomial(AbstractKernel):
 
         self.degree = degree
 
-        if isinstance(shift, nnx.Variable):
-            self.shift = shift
-        else:
-            self.shift = PositiveReal(shift)
-            if tp.TYPE_CHECKING:
-                self.shift = tp.cast(PositiveReal[ScalarArray], self.shift)
+        self.shift = shift
+        if tp.TYPE_CHECKING and not isinstance(shift, nnx.Variable):
+            self.shift = tp.cast(PositiveReal[ScalarArray], self.shift)
 
         if isinstance(variance, nnx.Variable):
             self.variance = variance
@@ -88,7 +85,9 @@ class Polynomial(AbstractKernel):
     def __call__(self, x: Float[Array, " D"], y: Float[Array, " D"]) -> ScalarFloat:
         x = self.slice_input(x)
         y = self.slice_input(y)
-        K = jnp.power(
-            self.shift.value + self.variance.value * jnp.dot(x, y), self.degree
+        shift_val = self.shift.value if hasattr(self.shift, "value") else self.shift
+        variance_val = (
+            self.variance.value if hasattr(self.variance, "value") else self.variance
         )
+        K = jnp.power(shift_val + variance_val * jnp.dot(x, y), self.degree)
         return K.squeeze()

gpjax/kernels/stationary/periodic.py

@@ -23,7 +23,6 @@ from gpjax.kernels.computations import (
     DenseKernelComputation,
 )
 from gpjax.kernels.stationary.base import StationaryKernel
-from gpjax.parameters import PositiveReal
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -72,10 +71,7 @@ class Periodic(StationaryKernel):
                 covariance matrix.
         """
 
-        if isinstance(period, nnx.Variable):
-            self.period = period
-        else:
-            self.period = PositiveReal(period)
+        self.period = period
 
         super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine)
 
@@ -84,8 +80,9 @@ class Periodic(StationaryKernel):
     ) -> Float[Array, ""]:
         x = self.slice_input(x)
         y = self.slice_input(y)
+        period_val = self.period.value if hasattr(self.period, "value") else self.period
         sine_squared = (
-            jnp.sin(jnp.pi * (x - y) / self.period.value) / self.lengthscale.value
+            jnp.sin(jnp.pi * (x - y) / period_val) / self.lengthscale.value
         ) ** 2
         K = self.variance.value * jnp.exp(-0.5 * jnp.sum(sine_squared, axis=0))
         return K.squeeze()

gpjax/kernels/stationary/powered_exponential.py

@@ -24,7 +24,6 @@ from gpjax.kernels.computations import (
 )
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import euclidean_distance
-from gpjax.parameters import SigmoidBounded
 from gpjax.typing import (
     Array,
     ScalarArray,
@@ -76,10 +75,7 @@ class PoweredExponential(StationaryKernel):
             compute_engine: the computation engine that the kernel uses to compute the
                 covariance matrix.
         """
-        if isinstance(power, nnx.Variable):
-            self.power = power
-        else:
-            self.power = SigmoidBounded(power)
+        self.power = power
 
         super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine)
 
@@ -88,7 +84,6 @@ class PoweredExponential(StationaryKernel):
     ) -> Float[Array, ""]:
         x = self.slice_input(x) / self.lengthscale.value
         y = self.slice_input(y) / self.lengthscale.value
-        K = self.variance.value * jnp.exp(
-            -(euclidean_distance(x, y) ** self.power.value)
-        )
+        power_val = self.power.value if hasattr(self.power, "value") else self.power
+        K = self.variance.value * jnp.exp(-(euclidean_distance(x, y) ** power_val))
         return K.squeeze()

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

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.12.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.12.0
+Version: 0.12.2
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -17,7 +17,7 @@ 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.10
 Requires-Dist: beartype>0.16.1
 Requires-Dist: flax>=0.10.0
 Requires-Dist: jax>=0.5.0
@@ -60,7 +60,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'
@@ -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
 
@@ -218,13 +209,14 @@ configuration in development mode.
 ```bash
 git clone https://github.com/JaxGaussianProcesses/GPJax.git
 cd GPJax
-uv sync --extra dev
+hatch env create
+hatch shell
 ```
 
 > We recommend you check your installation passes the supplied unit tests:
 >
 > ```python
-> uv run pytest --beartype-packages='gpjax'
+> hatch run dev:test
 > ```
 
 # Citing GPJax

{gpjax-0.12.0.dist-info → gpjax-0.12.2.dist-info}/RECORD

@@ -1,52 +1,52 @@
-gpjax/__init__.py,sha256=FSrKDFSQ7xDqwQGBWwEPqqjvYxEbhUPPestKLoAPjWA,1686
+gpjax/__init__.py,sha256=RzwpixFXn6HNHLVLy4LVXhFUk2c-_ce6n1gjZ2B93F0,1641
 gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
 gpjax/distributions.py,sha256=iKmeQ_NN2CIjRiuOeJlwEGASzGROi4ZCerVi1uY7zRM,7758
-gpjax/fit.py,sha256=R4TIPvBNHYSg9vBVp6is_QYENldRLIU_FklGE85C-aA,15046
-gpjax/gps.py,sha256=-Log0pcU8qmB5fUxfzoNjD0S64gpiypAjFzjGXX6w7I,30301
+gpjax/fit.py,sha256=I2sJVuKZii_d7MEcelHIivfM8ExYGMgdBuKKOT7Dw-A,15326
+gpjax/gps.py,sha256=ipaeYMnPffhKK_JsEHe4fF8GmolQIjXB1YbyfUIL8H4,30118
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
-gpjax/likelihoods.py,sha256=99oTZoWld1M7vxgGM0pNY5Hnt2Ajd2lQNqawzrLmwtk,9308
-gpjax/mean_functions.py,sha256=-sVYO1_LWE8f34rllUOuaT5sgGGAdxo99v5kRo2d4oM,6490
+gpjax/likelihoods.py,sha256=xwnSQpn6Aa-FPpEoDn_3xpBdPQAmHP97jP-9iJmT4G8,9087
+gpjax/mean_functions.py,sha256=KiHQXI-b7o0Vi5KQxGm6RNsUjitJc9jEOCq2GrSx4II,6531
 gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
-gpjax/objectives.py,sha256=Tm36h8fz_nWkZPlufMQzZWKK1ytrtT9yvvP8YdxYKNw,15359
-gpjax/parameters.py,sha256=qIEqyMKNd2n2Ak15PisCmqhX5qhsoRgng_s4doL96rE,7044
+gpjax/objectives.py,sha256=GvKbDIPqYjsc9FpiTccmZwRdHr2lCykgfxI9BX9I_GA,15362
+gpjax/parameters.py,sha256=hnyIKr6uIzd7Kb3KZC9WowR88ruQwUvdcto3cx2ZDv4,6756
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=rE3LarwIAkvDvLlWrz8Ww6BUBz88YHdV4ceY97r3IBw,28637
+gpjax/variational_families.py,sha256=TJGGkwkE805X4PQb-C32FxvD9B_OsFLWf6I-ZZvOUWk,29628
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
-gpjax/kernels/base.py,sha256=hOUXwarspDFnuI2_QreyIVPdz2fzRVJj4p3Zdu1touw,11606
+gpjax/kernels/base.py,sha256=4Lx8y3kPX4WqQZGRGAsBkqn_i6FlfoAhSn9Tv415xuQ,11551
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
-gpjax/kernels/approximations/rff.py,sha256=VbitjNuahFE5_IvCj1A0SxHhJXU0O0Qq0FMMVq8xA3E,4125
+gpjax/kernels/approximations/rff.py,sha256=GbNUmDPEKEKuMwxUcocxl_9IFR3Q9KEPZXzjy_ZD-2w,4043
 gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
 gpjax/kernels/computations/base.py,sha256=L6K0roxZbrYeJKxEw-yaTiK9Mtcv0YtZfWI2Xnau7i8,3616
-gpjax/kernels/computations/basis_functions.py,sha256=MPSo40NEx_ngnSLTa9ntVJzma_jugvm5dMpZd5MtG5M,2490
+gpjax/kernels/computations/basis_functions.py,sha256=_SFv4Tiwne40bxr1uVYpEjjZgjIQHKseLmss2Zgl1L4,2484
 gpjax/kernels/computations/constant_diagonal.py,sha256=JkQhLj7cK48IhOER4ivkALNhD1oQleKe-Rr9BtUJ6es,1984
 gpjax/kernels/computations/dense.py,sha256=vnW6XKQe4_gzpXRWTctxhgMA9-9TebdtiXzAqh_-j6g,1392
 gpjax/kernels/computations/diagonal.py,sha256=k1KqW0DwWRIBvbb7jzcKktXRfhXbcos3ncWrFplJ4W0,1768
-gpjax/kernels/computations/eigen.py,sha256=w7I7LK42j0ouchHCI1ltXx0lpwqvK1bRb4HclnF3rKs,1936
+gpjax/kernels/computations/eigen.py,sha256=NTHm-cn-RepYuXFrvXo2ih7Gtu1YR_pAg4Jb7IhE_o8,1930
 gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3J9_fWz1NlY,790
-gpjax/kernels/non_euclidean/graph.py,sha256=K4WIdX-dx1SsWuNHZnNjHFw8ElKZxGcReUiA3w4aCOI,4204
+gpjax/kernels/non_euclidean/graph.py,sha256=xTrx6ro8ubRXgM7Wgg6NmOyyEjEcGhzydY7KXueknCc,4120
 gpjax/kernels/non_euclidean/utils.py,sha256=z42aw8ga0zuREzHawemR9okttgrAUPmq-aN5HMt4SuY,1578
 gpjax/kernels/nonstationary/__init__.py,sha256=YpWQfOy_cqOKc5ezn37vqoK3Z6jznYiJz28BD_8F7AY,930
-gpjax/kernels/nonstationary/arccosine.py,sha256=2WV6aM0Z3-xXZnoPw-77n2CW62n-AZuJy-7AQ9xrMco,5858
+gpjax/kernels/nonstationary/arccosine.py,sha256=cqb8sqaNwW3fEbrA7MY9OF2KJFTkxHhqwmQtABE3G8w,5408
 gpjax/kernels/nonstationary/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
-gpjax/kernels/nonstationary/polynomial.py,sha256=arP8DK0jnBOaayDWcFvHF0pdu9FVhwzXdqjnHUAL2VI,3293
+gpjax/kernels/nonstationary/polynomial.py,sha256=CKc02C7Utgo-hhcOOCcKLdln5lj4vud_8M-JY7SevJ8,3388
 gpjax/kernels/stationary/__init__.py,sha256=j4BMTaQlIx2kNAT1Dkf4iO2rm-f7_oSVWNrk1bN0tqE,1406
 gpjax/kernels/stationary/base.py,sha256=25qDqpZP4gNtzbyzDCW-6u7rJfMqkg0dW88XUmTTupU,7078
 gpjax/kernels/stationary/matern12.py,sha256=DGjqw6VveYsyy0TrufyJJvCei7p9slnm2f0TgRGG7_U,1773
 gpjax/kernels/stationary/matern32.py,sha256=laLsJWJozJzpYHBzlkPUq0rWxz1eWEwGC36P2nPJuaQ,1966
 gpjax/kernels/stationary/matern52.py,sha256=VSByD2sb7k-DzRFjaz31P3Rtc4bPPhHvMshrxZNFnns,2019
-gpjax/kernels/stationary/periodic.py,sha256=IAbCxURtJEHGdmYzbdrsqRZ3zJ8F8tGQF9O7sggafZk,3598
-gpjax/kernels/stationary/powered_exponential.py,sha256=8qT91IWKJK7PpEtFcX4MVu1ahWMOFOZierPko4JCjKA,3776
-gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UOl8uOFT3lCutleSvo4,3496
+gpjax/kernels/stationary/periodic.py,sha256=f4PhWhKg-pJsEBGzEMK9pdbylO84GPKhzHlBC83ZVWw,3528
+gpjax/kernels/stationary/powered_exponential.py,sha256=xuFGuIK0mKNMU3iLtZMXZTHXJuMFAMoX7gAtXefCdqU,3679
+gpjax/kernels/stationary/rational_quadratic.py,sha256=zHo2LVW65T52XET4Hx9JaKO0TfxylV8WRUtP7sUUOx0,3418
 gpjax/kernels/stationary/rbf.py,sha256=euHUs6FdfRICQcabAWE4MX-7GEDr2TxgZWdFQiXr9Bw,1690
 gpjax/kernels/stationary/utils.py,sha256=6BI9EBcCzeeKx-XH-MfW1ORmtU__tPX5zyvfLhpkBsU,2180
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
 gpjax/linalg/__init__.py,sha256=F8mxk_9Zc2nFd7Q-unjJ50_6rXEKzZj572WsU_jUKqI,547
 gpjax/linalg/operations.py,sha256=xvhOy5P4FmUCPWjIVNdg1yDXaoFQ48anFUfR-Tnfr6k,6480
 gpjax/linalg/operators.py,sha256=arxRGwcoAy_RqUYqBpZ3XG6OXbjShUl7m8sTpg85npE,11608
-gpjax/linalg/utils.py,sha256=DGX40TDhmfYn7JBxElpBm_9W0cetm0HZUK7B3j74xxo,895
-gpjax-0.12.0.dist-info/METADATA,sha256=8lLQb5SUvWvniry-zBOR3wzm03tXvHe7Lzry_Ho3peE,10562
-gpjax-0.12.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.12.0.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.12.0.dist-info/RECORD,,
+gpjax/linalg/utils.py,sha256=fKV8G_iKZVhNkNvN20D_dQEi93-8xosGbXBP-v7UEyo,2020
+gpjax-0.12.2.dist-info/METADATA,sha256=eckQKXiBXi8XbBeJFviBAIPdBGVWGFQg7wVZwMfPPxs,10129
+gpjax-0.12.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.12.2.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.12.2.dist-info/RECORD,,