gpjax 0.12.0__py3-none-any.whl → 0.13.0__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.13.0"
 
 __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

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

@@ -1,13 +1,13 @@
 """Compute Random Fourier Feature (RFF) kernel approximations."""
 
 import beartype.typing as tp
+from flax import nnx
 import jax.random as jr
 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,
@@ -55,7 +55,7 @@ class RFF(AbstractKernel):
         self._check_valid_base_kernel(base_kernel)
         self.base_kernel = base_kernel
         self.num_basis_fns = num_basis_fns
-        self.frequencies = frequencies
+        self.frequencies = nnx.data(frequencies)
         self.compute_engine = compute_engine
 
         if self.frequencies is None:
@@ -66,10 +66,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):
@@ -256,7 +253,7 @@ class CombinationKernel(AbstractKernel):
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),
     ):
         # Add kernels to a list, flattening out instances of this class therein, as in GPFlow kernels.
-        kernels_list: list[AbstractKernel] = []
+        kernels_list: list[AbstractKernel] = nnx.List([])
         for kernel in kernels:
             if not isinstance(kernel, AbstractKernel):
                 raise TypeError("can only combine Kernel instances")  # pragma: no cover

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/base.py

@@ -127,7 +127,7 @@ def _check_lengthscale_dims_compat(
     """
 
     if isinstance(lengthscale, nnx.Variable):
-        return _check_lengthscale_dims_compat_old(lengthscale.value, n_dims)
+        return _check_lengthscale_dims_compat(lengthscale.value, n_dims)
 
     lengthscale = jnp.asarray(lengthscale)
     ls_shape = jnp.shape(lengthscale)
@@ -146,35 +146,6 @@ def _check_lengthscale_dims_compat(
         return n_dims
 
 
-def _check_lengthscale_dims_compat_old(
-    lengthscale: tp.Union[LengthscaleCompatible, nnx.Variable[Lengthscale]],
-    n_dims: tp.Union[int, None],
-):
-    r"""Check that the lengthscale is compatible with n_dims.
-
-    If possible, infer the number of input dimensions from the lengthscale.
-    """
-
-    if isinstance(lengthscale, nnx.Variable):
-        return _check_lengthscale_dims_compat_old(lengthscale.value, n_dims)
-
-    lengthscale = jnp.asarray(lengthscale)
-    ls_shape = jnp.shape(lengthscale)
-
-    if ls_shape == ():
-        return lengthscale, n_dims
-    elif ls_shape != () and n_dims is None:
-        return lengthscale, ls_shape[0]
-    elif ls_shape != () and n_dims is not None:
-        if ls_shape != (n_dims,):
-            raise ValueError(
-                "Expected `lengthscale` to be compatible with the number "
-                f"of input dimensions. Got `lengthscale` with shape {ls_shape}, "
-                f"but the number of input dimensions is {n_dims}."
-            )
-        return lengthscale, n_dims
-
-
 def _check_lengthscale(lengthscale: tp.Any):
     """Check that the lengthscale is a valid value."""
 

gpjax/kernels/stationary/matern12.py

@@ -35,7 +35,7 @@ class Matern12(StationaryKernel):
     lengthscale parameter $\ell$ and variance $\sigma^2$.
 
     $$
-    k(x, y) = \sigma^2\exp\Bigg(-\frac{\lvert x-y \rvert}{2\ell^2}\Bigg)
+    k(x, y) = \sigma^2\exp\Bigg(-\frac{\lvert x-y \rvert}{2\ell}\Bigg)
     $$
     """
 

gpjax/kernels/stationary/matern32.py

@@ -32,7 +32,7 @@ class Matern32(StationaryKernel):
     lengthscale parameter $\ell$ and variance $\sigma^2$.
 
     $$
-    k(x, y) = \sigma^2 \exp \Bigg(1+ \frac{\sqrt{3}\lvert x-y \rvert}{\ell^2} \ \Bigg)\exp\Bigg(-\frac{\sqrt{3}\lvert x-y\rvert}{\ell^2} \Bigg)
+    k(x, y) = \sigma^2 \exp \Bigg(1+ \frac{\sqrt{3}\lvert x-y \rvert}{\ell} \ \Bigg)\exp\Bigg(-\frac{\sqrt{3}\lvert x-y\rvert}{\ell^2} \Bigg)
     $$
     """
 

gpjax/kernels/stationary/matern52.py

@@ -33,7 +33,7 @@ class Matern52(StationaryKernel):
     lengthscale parameter $\ell$ and variance $\sigma^2$.
 
     $$
-    k(x, y) = \sigma^2 \exp \Bigg(1+ \frac{\sqrt{5}\lvert x-y \rvert}{\ell^2} + \frac{5\lvert x - y \rvert^2}{3\ell^2} \Bigg)\exp\Bigg(-\frac{\sqrt{5}\lvert x-y\rvert}{\ell^2} \Bigg)
+    k(x, y) = \sigma^2 \exp \Bigg(1+ \frac{\sqrt{5}\lvert x-y \rvert}{\ell} + \frac{5\lvert x - y \rvert^2}{3\ell^2} \Bigg)\exp\Bigg(-\frac{\sqrt{5}\lvert x-y\rvert}{\ell^2} \Bigg)
     $$
     """
 

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