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

gpjax/__init__.py

@@ -40,7 +40,7 @@ __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.11.1"
+__version__ = "0.12.0"
 
 __all__ = [
     "base",

gpjax/citation.py

@@ -8,7 +8,12 @@ from beartype.typing import (
     Dict,
     Union,
 )
-from jaxlib.xla_extension import PjitFunction
+
+try:
+    # safely removable once jax>=0.6.0
+    from jaxlib.xla_extension import PjitFunction
+except ModuleNotFoundError:
+    from jaxlib._jax import PjitFunction
 
 from gpjax.kernels import (
     RFF,
@@ -45,7 +50,7 @@ class AbstractCitation:
 
 
 class NullCitation(AbstractCitation):
-    def __str__(self) -> str:
+    def as_str(self) -> str:
         return (
             "No citation available. If you think this is an error, please open a pull"
             " request."

gpjax/distributions.py

@@ -17,11 +17,6 @@
 from beartype.typing import (
     Optional,
 )
-import cola
-from cola.linalg.decompositions import Cholesky
-from cola.ops import (
-    LinearOperator,
-)
 from jax import vmap
 import jax.numpy as jnp
 import jax.random as jr
@@ -30,7 +25,14 @@ from numpyro.distributions import constraints
 from numpyro.distributions.distribution import Distribution
 from numpyro.distributions.util import is_prng_key
 
-from gpjax.lower_cholesky import lower_cholesky
+from gpjax.linalg.operations import (
+    diag,
+    logdet,
+    lower_cholesky,
+    solve,
+)
+from gpjax.linalg.operators import LinearOperator
+from gpjax.linalg.utils import psd
 from gpjax.typing import (
     Array,
     ScalarFloat,
@@ -47,7 +49,7 @@ class GaussianDistribution(Distribution):
         validate_args=None,
     ):
         self.loc = loc
-        self.scale = cola.PSD(scale)
+        self.scale = psd(scale)
         batch_shape = ()
         event_shape = jnp.shape(self.loc)
         super().__init__(batch_shape, event_shape, validate_args=validate_args)
@@ -76,13 +78,12 @@ class GaussianDistribution(Distribution):
     @property
     def variance(self) -> Float[Array, " N"]:
         r"""Calculates the variance."""
-        return cola.diag(self.scale)
+        return diag(self.scale)
 
     def entropy(self) -> ScalarFloat:
         r"""Calculates the entropy of the distribution."""
         return 0.5 * (
-            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi))
-            + cola.logdet(self.scale, Cholesky(), Cholesky())
+            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi)) + logdet(self.scale)
         )
 
     def median(self) -> Float[Array, " N"]:
@@ -104,7 +105,7 @@ class GaussianDistribution(Distribution):
 
     def stddev(self) -> Float[Array, " N"]:
         r"""Calculates the standard deviation."""
-        return jnp.sqrt(cola.diag(self.scale))
+        return jnp.sqrt(diag(self.scale))
 
     #     @property
     #     def event_shape(self) -> Tuple:
@@ -129,9 +130,7 @@ class GaussianDistribution(Distribution):
 
         # compute the pdf, -1/2[ n log(2π) + log|Σ| + (y - µ)ᵀΣ⁻¹(y - µ) ]
         return -0.5 * (
-            n * jnp.log(2.0 * jnp.pi)
-            + cola.logdet(sigma, Cholesky(), Cholesky())
-            + diff.T @ cola.solve(sigma, diff, Cholesky())
+            n * jnp.log(2.0 * jnp.pi) + logdet(sigma) + diff.T @ solve(sigma, diff)
         )
 
     #     def _sample_n(self, key: KeyArray, n: int) -> Float[Array, "n N"]:
@@ -219,53 +218,14 @@ def _kl_divergence(q: GaussianDistribution, p: GaussianDistribution) -> ScalarFl
 
     # trace term, tr[Σp⁻¹ Σq] = tr[(LpLpᵀ)⁻¹(LqLqᵀ)] = tr[(Lp⁻¹Lq)(Lp⁻¹Lq)ᵀ] = (fr[LqLp⁻¹])²
     trace = _frobenius_norm_squared(
-        cola.solve(sqrt_p, sqrt_q.to_dense(), Cholesky())
+        solve(sqrt_p, sqrt_q.to_dense())
     )  # TODO: Not most efficient, given the `to_dense()` call (e.g., consider diagonal p and q). Need to abstract solving linear operator against another linear operator.
 
     # Mahalanobis term, (μp - μq)ᵀ Σp⁻¹ (μp - μq) = tr [(μp - μq)ᵀ [LpLpᵀ]⁻¹ (μp - μq)] = (fr[Lp⁻¹(μp - μq)])²
-    mahalanobis = jnp.sum(jnp.square(cola.solve(sqrt_p, diff, Cholesky())))
+    mahalanobis = jnp.sum(jnp.square(solve(sqrt_p, diff)))
 
     # KL[q(x)||p(x)] = [ [(μp - μq)ᵀ Σp⁻¹ (μp - μq)] - n - log|Σq| + log|Σp| + tr[Σp⁻¹ Σq] ] / 2
-    return (
-        mahalanobis
-        - n_dim
-        - cola.logdet(sigma_q, Cholesky(), Cholesky())
-        + cola.logdet(sigma_p, Cholesky(), Cholesky())
-        + trace
-    ) / 2.0
-
-
-# def _check_loc_scale(loc: Optional[Any], scale: Optional[Any]) -> None:
-#     r"""Checks that the inputs are correct."""
-#     if loc is None and scale is None:
-#         raise ValueError("At least one of `loc` or `scale` must be specified.")
-
-#     if loc is not None and loc.ndim < 1:
-#         raise ValueError("The parameter `loc` must have at least one dimension.")
-
-#     if scale is not None and len(scale.shape) < 2:  # scale.ndim < 2:
-#         raise ValueError(
-#             "The `scale` must have at least two dimensions, but "
-#             f"`scale.shape = {scale.shape}`."
-#         )
-
-#     if scale is not None and not isinstance(scale, LinearOperator):
-#         raise ValueError(
-#             f"The `scale` must be a CoLA LinearOperator but got {type(scale)}"
-#         )
-
-#     if scale is not None and (scale.shape[-1] != scale.shape[-2]):
-#         raise ValueError(
-#             f"The `scale` must be a square matrix, but `scale.shape = {scale.shape}`."
-#         )
-
-#     if loc is not None:
-#         num_dims = loc.shape[-1]
-#         if scale is not None and (scale.shape[-1] != num_dims):
-#             raise ValueError(
-#                 f"Shapes are not compatible: `loc.shape = {loc.shape}` and "
-#                 f"`scale.shape = {scale.shape}`."
-#             )
+    return (mahalanobis - n_dim - logdet(sigma_q) + logdet(sigma_p) + trace) / 2.0
 
 
 __all__ = [

gpjax/fit.py

@@ -15,13 +15,13 @@
 
 import typing as tp
 
+from flax import nnx
 import jax
+from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random as jr
-import optax as ox
-from flax import nnx
-from jax.flatten_util import ravel_pytree
 from numpyro.distributions.transforms import Transform
+import optax as ox
 from scipy.optimize import minimize
 
 from gpjax.dataset import Dataset

gpjax/gps.py

@@ -16,11 +16,6 @@
 from abc import abstractmethod
 
 import beartype.typing as tp
-from cola.annotations import PSD
-from cola.linalg.algorithm_base import Algorithm
-from cola.linalg.decompositions.decompositions import Cholesky
-from cola.linalg.inverse.inv import solve
-from cola.ops.operators import I_like
 from flax import nnx
 import jax.numpy as jnp
 import jax.random as jr
@@ -38,7 +33,13 @@ from gpjax.likelihoods import (
     Gaussian,
     NonGaussian,
 )
-from gpjax.lower_cholesky import lower_cholesky
+from gpjax.linalg import (
+    Dense,
+    Identity,
+    psd,
+    solve,
+)
+from gpjax.linalg.operations import lower_cholesky
 from gpjax.mean_functions import AbstractMeanFunction
 from gpjax.parameters import (
     Parameter,
@@ -251,8 +252,8 @@ class Prior(AbstractPrior[M, K]):
         x = test_inputs
         mx = self.mean_function(x)
         Kxx = self.kernel.gram(x)
-        Kxx += I_like(Kxx) * self.jitter
-        Kxx = PSD(Kxx)
+        Kxx_dense = Kxx.to_dense() + Identity(Kxx.shape).to_dense() * self.jitter
+        Kxx = psd(Dense(Kxx_dense))
 
         return GaussianDistribution(jnp.atleast_1d(mx.squeeze()), Kxx)
 
@@ -315,15 +316,13 @@ class Prior(AbstractPrior[M, K]):
         if (not isinstance(num_samples, int)) or num_samples <= 0:
             raise ValueError("num_samples must be a positive integer")
 
-        # sample fourier features
         fourier_feature_fn = _build_fourier_features_fn(self, num_features, key)
 
-        # sample fourier weights
-        feature_weights = jr.normal(key, [num_samples, 2 * num_features])  # [B, L]
+        feature_weights = jr.normal(key, [num_samples, 2 * num_features])
 
         def sample_fn(test_inputs: Float[Array, "N D"]) -> Float[Array, "N B"]:
-            feature_evals = fourier_feature_fn(test_inputs)  # [N, L]
-            evaluated_sample = jnp.inner(feature_evals, feature_weights)  # [N, B]
+            feature_evals = fourier_feature_fn(test_inputs)
+            evaluated_sample = jnp.inner(feature_evals, feature_weights)
             return self.mean_function(test_inputs) + evaluated_sample
 
         return sample_fn
@@ -504,24 +503,23 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
 
         # Precompute Gram matrix, Kxx, at training inputs, x
         Kxx = self.prior.kernel.gram(x)
-        Kxx += I_like(Kxx) * self.jitter
+        Kxx_dense = Kxx.to_dense() + Identity(Kxx.shape).to_dense() * self.jitter
+        Kxx = Dense(Kxx_dense)
 
-        # Σ = Kxx + Io²
-        Sigma = Kxx + I_like(Kxx) * obs_noise
-        Sigma = PSD(Sigma)
+        Sigma_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * obs_noise
+        Sigma = psd(Dense(Sigma_dense))
 
         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, Cholesky())
+        Sigma_inv_Kxt = solve(Sigma, Kxt)
 
-        # μt  +  Ktx (Kxx + Io²)⁻¹ (y  -  μx)
         mean = mean_t + jnp.matmul(Sigma_inv_Kxt.T, y - mx)
 
         # Ktt  -  Ktx (Kxx + Io²)⁻¹ Kxt, TODO: Take advantage of covariance structure to compute Schur complement more efficiently.
-        covariance = Ktt - jnp.matmul(Kxt.T, Sigma_inv_Kxt)
-        covariance += I_like(covariance) * self.prior.jitter
-        covariance = PSD(covariance)
+        covariance = Ktt.to_dense() - jnp.matmul(Kxt.T, Sigma_inv_Kxt)
+        covariance += jnp.eye(covariance.shape[0]) * self.prior.jitter
+        covariance = psd(Dense(covariance))
 
         return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), covariance)
 
@@ -531,7 +529,6 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
         train_data: Dataset,
         key: KeyArray,
         num_features: int | None = 100,
-        solver_algorithm: tp.Optional[Algorithm] = Cholesky(),
     ) -> FunctionalSample:
         r"""Draw approximate samples from the Gaussian process posterior.
 
@@ -565,11 +562,6 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
             key (KeyArray): The random seed used for the sample(s).
             num_features (int): The number of features used when approximating the
                 kernel.
-            solver_algorithm (Optional[Algorithm], optional): The algorithm to use for the solves of
-                the inverse of the covariance matrix. See the
-                [CoLA documentation](https://cola.readthedocs.io/en/latest/package/cola.linalg.html#algorithms)
-                for which solver to pick. For PSD matrices, CoLA currently recommends Cholesky() for small
-                matrices and CG() for larger matrices. Select Auto() to let CoLA decide. Defaults to Cholesky().
 
         Returns:
             FunctionalSample: A function representing an approximate sample from the Gaussian
@@ -581,31 +573,25 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
         # sample fourier features
         fourier_feature_fn = _build_fourier_features_fn(self.prior, num_features, key)
 
-        # sample fourier weights
-        fourier_weights = jr.normal(key, [num_samples, 2 * num_features])  # [B, L]
+        fourier_weights = jr.normal(key, [num_samples, 2 * num_features])
 
-        # sample weights v for canonical features
-        # v = Σ⁻¹ (y + ε - ɸ⍵) for  Σ = Kxx + Io² and ε ᯈ N(0, o²)
         obs_var = self.likelihood.obs_stddev.value**2
-        Kxx = self.prior.kernel.gram(train_data.X)  #  [N, N]
-        Sigma = Kxx + I_like(Kxx) * (obs_var + self.jitter)  #  [N, N]
-        eps = jnp.sqrt(obs_var) * jr.normal(key, [train_data.n, num_samples])  #  [N, B]
-        y = train_data.y - self.prior.mean_function(train_data.X)  # account for mean
+        Kxx = self.prior.kernel.gram(train_data.X)
+        Sigma = Kxx + jnp.eye(Kxx.shape[0]) * (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)
         canonical_weights = solve(
             Sigma,
             y + eps - jnp.inner(Phi, fourier_weights),
-            solver_algorithm,
         )  #  [N, B]
 
         def sample_fn(test_inputs: Float[Array, "n D"]) -> Float[Array, "n B"]:
-            fourier_features = fourier_feature_fn(test_inputs)  # [n, L]
-            weight_space_contribution = jnp.inner(
-                fourier_features, fourier_weights
-            )  # [n, B]
+            fourier_features = fourier_feature_fn(test_inputs)
+            weight_space_contribution = jnp.inner(fourier_features, fourier_weights)
             canonical_features = self.prior.kernel.cross_covariance(
                 test_inputs, train_data.X
-            )  # [n, N]
+            )
             function_space_contribution = jnp.matmul(
                 canonical_features, canonical_weights
             )
@@ -689,8 +675,8 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
 
         # Precompute lower triangular of Gram matrix, Lx, at training inputs, x
         Kxx = kernel.gram(x)
-        Kxx += I_like(Kxx) * self.prior.jitter
-        Kxx = PSD(Kxx)
+        Kxx_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * self.prior.jitter
+        Kxx = psd(Dense(Kxx_dense))
         Lx = lower_cholesky(Kxx)
 
         # Unpack test inputs
@@ -702,7 +688,7 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         mean_t = mean_function(t)
 
         # Lx⁻¹ Kxt
-        Lx_inv_Kxt = solve(Lx, Ktx.T, Cholesky())
+        Lx_inv_Kxt = solve(Lx, Ktx.T)
 
         # Whitened function values, wx, corresponding to the inputs, x
         wx = self.latent.value
@@ -711,9 +697,9 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         mean = mean_t + jnp.matmul(Lx_inv_Kxt.T, wx)
 
         # Ktt - Ktx Kxx⁻¹ Kxt, TODO: Take advantage of covariance structure to compute Schur complement more efficiently.
-        covariance = Ktt - jnp.matmul(Lx_inv_Kxt.T, Lx_inv_Kxt)
-        covariance += I_like(covariance) * self.prior.jitter
-        covariance = PSD(covariance)
+        covariance = Ktt.to_dense() - jnp.matmul(Lx_inv_Kxt.T, Lx_inv_Kxt)
+        covariance += jnp.eye(covariance.shape[0]) * self.prior.jitter
+        covariance = psd(Dense(covariance))
 
         return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), covariance)
 

gpjax/kernels/base.py

@@ -17,7 +17,6 @@ import abc
 import functools as ft
 
 import beartype.typing as tp
-from cola.ops.operator_base import LinearOperator
 from flax import nnx
 import jax.numpy as jnp
 from jaxtyping import (
@@ -29,6 +28,7 @@ from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
 )
+from gpjax.linalg import LinearOperator
 from gpjax.parameters import (
     Parameter,
     Real,

gpjax/kernels/computations/base.py

@@ -16,11 +16,6 @@
 import abc
 import typing as tp
 
-from cola.annotations import PSD
-from cola.ops.operators import (
-    Dense,
-    Diagonal,
-)
 from jax import vmap
 from jaxtyping import (
     Float,
@@ -28,6 +23,11 @@ from jaxtyping import (
 )
 
 import gpjax
+from gpjax.linalg import (
+    Dense,
+    Diagonal,
+    psd,
+)
 from gpjax.typing import Array
 
 K = tp.TypeVar("K", bound="gpjax.kernels.base.AbstractKernel")  # noqa: F821
@@ -69,7 +69,7 @@ class AbstractKernelComputation:
             The Gram covariance of the kernel function as a linear operator.
         """
         Kxx = self.cross_covariance(kernel, x, x)
-        return PSD(Dense(Kxx))
+        return psd(Dense(Kxx))
 
     @abc.abstractmethod
     def _cross_covariance(
@@ -93,7 +93,7 @@ class AbstractKernelComputation:
         return self._cross_covariance(kernel, x, y)
 
     def _diagonal(self, kernel: K, inputs: Num[Array, "N D"]) -> Diagonal:
-        return PSD(Diagonal(diag=vmap(lambda x: kernel(x, x))(inputs)))
+        return psd(Diagonal(vmap(lambda x: kernel(x, x))(inputs)))
 
     def diagonal(self, kernel: K, inputs: Num[Array, "N D"]) -> Diagonal:
         r"""For a given kernel, compute the elementwise diagonal of the

gpjax/kernels/computations/basis_functions.py

@@ -1,18 +1,19 @@
 import typing as tp
 
-from cola.annotations import PSD
-from cola.ops.operators import Dense
 import jax.numpy as jnp
 from jaxtyping import Float
 
 import gpjax
 from gpjax.kernels.computations.base import AbstractKernelComputation
+from gpjax.linalg import (
+    Dense,
+    Diagonal,
+    psd,
+)
 from gpjax.typing import Array
 
 K = tp.TypeVar("K", bound="gpjax.kernels.approximations.RFF")  # noqa: F821
 
-from cola.ops import Diagonal
-
 # TODO: Use low rank linear operator!
 
 
@@ -28,7 +29,7 @@ class BasisFunctionComputation(AbstractKernelComputation):
 
     def _gram(self, kernel: K, inputs: Float[Array, "N D"]) -> Dense:
         z1 = self.compute_features(kernel, inputs)
-        return PSD(Dense(self.scaling(kernel) * jnp.matmul(z1, z1.T)))
+        return psd(Dense(self.scaling(kernel) * jnp.matmul(z1, z1.T)))
 
     def diagonal(self, kernel: K, inputs: Float[Array, "N D"]) -> Diagonal:
         r"""For a given kernel, compute the elementwise diagonal of the

gpjax/kernels/computations/constant_diagonal.py

@@ -15,36 +15,34 @@
 
 import typing as tp
 
-from cola.annotations import PSD
-from cola.ops.operators import (
-    Diagonal,
-    Identity,
-    Product,
-)
 from jax import vmap
 import jax.numpy as jnp
 from jaxtyping import Float
 
 import gpjax
 from gpjax.kernels.computations import AbstractKernelComputation
+from gpjax.linalg import (
+    Diagonal,
+    psd,
+)
 from gpjax.typing import Array
 
 K = tp.TypeVar("K", bound="gpjax.kernels.base.AbstractKernel")  # noqa: F821
-ConstantDiagonalType = Product
+ConstantDiagonalType = Diagonal
 
 
 class ConstantDiagonalKernelComputation(AbstractKernelComputation):
     r"""Computation engine for constant diagonal kernels."""
 
-    def gram(self, kernel: K, x: Float[Array, "N D"]) -> Product:
+    def gram(self, kernel: K, x: Float[Array, "N D"]) -> Diagonal:
         value = kernel(x[0], x[0])
-        dtype = value.dtype
-        shape = (x.shape[0], x.shape[0])
-        return PSD(jnp.atleast_1d(value) * Identity(shape=shape, dtype=dtype))
+        # Create a diagonal matrix with constant values
+        diag = jnp.full(x.shape[0], value)
+        return psd(Diagonal(diag))
 
     def _diagonal(self, kernel: K, inputs: Float[Array, "N D"]) -> Diagonal:
         diag = vmap(lambda x: kernel(x, x))(inputs)
-        return PSD(Diagonal(diag=diag))
+        return psd(Diagonal(diag))
 
     def _cross_covariance(
         self, kernel: K, x: Float[Array, "N D"], y: Float[Array, "M D"]

gpjax/kernels/computations/diagonal.py

@@ -14,16 +14,16 @@
 # ==============================================================================
 
 import beartype.typing as tp
-from cola.annotations import PSD
-from cola.ops.operators import (
-    Diagonal,
-    LinearOperator,
-)
 from jax import vmap
 from jaxtyping import Float
 
 import gpjax  # noqa: F401
 from gpjax.kernels.computations import AbstractKernelComputation
+from gpjax.linalg import (
+    Diagonal,
+    LinearOperator,
+    psd,
+)
 from gpjax.typing import Array
 
 Kernel = tp.TypeVar("Kernel", bound="gpjax.kernels.base.AbstractKernel")  # noqa: F821
@@ -35,7 +35,7 @@ class DiagonalKernelComputation(AbstractKernelComputation):
     """
 
     def gram(self, kernel: Kernel, x: Float[Array, "N D"]) -> LinearOperator:
-        return PSD(Diagonal(diag=vmap(lambda x: kernel(x, x))(x)))
+        return psd(Diagonal(vmap(lambda x: kernel(x, x))(x)))
 
     def _cross_covariance(
         self, kernel: Kernel, x: Float[Array, "N D"], y: Float[Array, "M D"]

gpjax/linalg/__init__.py

@@ -0,0 +1,37 @@
+"""Linear algebra module for GPJax."""
+
+from gpjax.linalg.operations import (
+    diag,
+    logdet,
+    lower_cholesky,
+    solve,
+)
+from gpjax.linalg.operators import (
+    BlockDiag,
+    Dense,
+    Diagonal,
+    Identity,
+    Kronecker,
+    LinearOperator,
+    Triangular,
+)
+from gpjax.linalg.utils import (
+    PSD,
+    psd,
+)
+
+__all__ = [
+    "LinearOperator",
+    "Dense",
+    "Diagonal",
+    "Identity",
+    "Triangular",
+    "BlockDiag",
+    "Kronecker",
+    "lower_cholesky",
+    "solve",
+    "logdet",
+    "diag",
+    "psd",
+    "PSD",
+]