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

gpjax/objectives.py

@@ -1,13 +1,5 @@
 from typing import TypeVar
 
-from cola.annotations import PSD
-from cola.linalg.decompositions.decompositions import Cholesky
-from cola.linalg.inverse.inv import (
-    inv,
-    solve,
-)
-from cola.linalg.trace.diag_trace import diag
-from cola.ops.operators import I_like
 from flax import nnx
 from jax import vmap
 import jax.numpy as jnp
@@ -22,7 +14,12 @@ from gpjax.gps import (
     ConjugatePosterior,
     NonConjugatePosterior,
 )
-from gpjax.lower_cholesky import lower_cholesky
+from gpjax.linalg import (
+    Dense,
+    lower_cholesky,
+    psd,
+    solve,
+)
 from gpjax.typing import (
     Array,
     ScalarFloat,
@@ -100,9 +97,9 @@ def conjugate_mll(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat:
 
     # Σ = (Kxx + Io²) = LLᵀ
     Kxx = posterior.prior.kernel.gram(x)
-    Kxx += I_like(Kxx) * posterior.prior.jitter
-    Sigma = Kxx + I_like(Kxx) * obs_noise
-    Sigma = PSD(Sigma)
+    Kxx_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * posterior.prior.jitter
+    Sigma_dense = Kxx_dense + jnp.eye(Kxx.shape[0]) * obs_noise
+    Sigma = psd(Dense(Sigma_dense))
 
     # p(y | x, θ), where θ are the model hyperparameters:
     mll = GaussianDistribution(jnp.atleast_1d(mx.squeeze()), Sigma)
@@ -164,11 +161,14 @@ def conjugate_loocv(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat
 
     # Σ = (Kxx + Io²)
     Kxx = posterior.prior.kernel.gram(x)
-    Sigma = Kxx + I_like(Kxx) * (obs_var + posterior.prior.jitter)
-    Sigma = PSD(Sigma)  # [N, N]
+    Sigma_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * (
+        obs_var + posterior.prior.jitter
+    )
+    Sigma = psd(Dense(Sigma_dense))  # [N, N]
 
-    Sigma_inv_y = solve(Sigma, y - mx, Cholesky())  # [N, 1]
-    Sigma_inv_diag = diag(inv(Sigma, Cholesky()))[:, None]  # [N, 1]
+    Sigma_inv_y = solve(Sigma, y - mx)  # [N, 1]
+    Sigma_inv = jnp.linalg.inv(Sigma.to_dense())
+    Sigma_inv_diag = jnp.diag(Sigma_inv)[:, None]  # [N, 1]
 
     loocv_means = mx + (y - mx) - Sigma_inv_y / Sigma_inv_diag
     loocv_stds = jnp.sqrt(1.0 / Sigma_inv_diag)
@@ -213,8 +213,8 @@ def log_posterior_density(
 
     # Gram matrix
     Kxx = posterior.prior.kernel.gram(x)
-    Kxx += I_like(Kxx) * posterior.prior.jitter
-    Kxx = PSD(Kxx)
+    Kxx_dense = Kxx.to_dense() + jnp.eye(Kxx.shape[0]) * posterior.prior.jitter
+    Kxx = psd(Dense(Kxx_dense))
     Lx = lower_cholesky(Kxx)
 
     # Compute the prior mean function
@@ -349,8 +349,8 @@ 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 += I_like(Kzz) * variational_family.jitter
-    Kzz = PSD(Kzz)
+    Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * variational_family.jitter
+    Kzz = psd(Dense(Kzz_dense))
     Kzx = kernel.cross_covariance(z, x)
     Kxx_diag = vmap(kernel, in_axes=(0, 0))(x, x)
     μx = mean_function(x)
@@ -383,7 +383,7 @@ def collapsed_elbo(variational_family: VF, data: Dataset) -> ScalarFloat:
     #
     #   with A and B defined as above.
 
-    A = solve(Lz, Kzx, Cholesky()) / jnp.sqrt(noise)
+    A = solve(Lz, Kzx) / jnp.sqrt(noise)
 
     # AAᵀ
     AAT = jnp.matmul(A, A.T)

gpjax/parameters.py

@@ -21,23 +21,20 @@ def transform(
     r"""Transforms parameters using a bijector.
 
     Example:
-    ```pycon
         >>> from gpjax.parameters import PositiveReal, transform
         >>> import jax.numpy as jnp
         >>> import numpyro.distributions.transforms as npt
         >>> from flax import nnx
         >>> params = nnx.State(
-        >>>     {
-        >>>         "a": PositiveReal(jnp.array([1.0])),
-        >>>         "b": PositiveReal(jnp.array([2.0])),
-        >>>     }
-        >>> )
+        ...     {
+        ...         "a": PositiveReal(jnp.array([1.0])),
+        ...         "b": PositiveReal(jnp.array([2.0])),
+        ...     }
+        ... )
         >>> params_bijection = {'positive': npt.SoftplusTransform()}
         >>> transformed_params = transform(params, params_bijection)
         >>> print(transformed_params["a"].value)
-         [1.3132617]
-    ```
-
+        [1.3132617]
 
     Args:
         params: A nnx.State object containing parameters to be transformed.
@@ -49,7 +46,7 @@ def transform(
     """
 
     def _inner(param):
-        bijector = params_bijection.get(param._tag, npt.IdentityTransform())
+        bijector = params_bijection.get(param.tag, npt.IdentityTransform())
         if inverse:
             transformed_value = bijector.inv(param.value)
         else:
@@ -60,10 +57,11 @@ def transform(
 
     gp_params, *other_params = params.split(Parameter, ...)
 
+    # Transform each parameter in the state
     transformed_gp_params: nnx.State = jtu.tree_map(
-        lambda x: _inner(x),
+        lambda x: _inner(x) if isinstance(x, Parameter) else x,
         gp_params,
-        is_leaf=lambda x: isinstance(x, nnx.VariableState),
+        is_leaf=lambda x: isinstance(x, Parameter),
     )
     return nnx.State.merge(transformed_gp_params, *other_params)
 
@@ -79,7 +77,7 @@ class Parameter(nnx.Variable[T]):
         _check_is_arraylike(value)
 
         super().__init__(value=jnp.asarray(value), **kwargs)
-        self._tag = tag
+        self.tag = tag
 
 
 class NonNegativeReal(Parameter[T]):

gpjax/variational_families.py

@@ -16,15 +16,6 @@
 import abc
 
 import beartype.typing as tp
-from cola.annotations import PSD
-from cola.linalg.decompositions.decompositions import Cholesky
-from cola.linalg.inverse.inv import solve
-from cola.ops.operators import (
-    Dense,
-    I_like,
-    Identity,
-    Triangular,
-)
 from flax import nnx
 import jax.numpy as jnp
 import jax.scipy as jsp
@@ -41,7 +32,14 @@ from gpjax.likelihoods import (
     Gaussian,
     NonGaussian,
 )
-from gpjax.lower_cholesky import lower_cholesky
+from gpjax.linalg import (
+    Dense,
+    Identity,
+    Triangular,
+    lower_cholesky,
+    psd,
+    solve,
+)
 from gpjax.mean_functions import AbstractMeanFunction
 from gpjax.parameters import (
     LowerTriangular,
@@ -189,7 +187,7 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz = PSD(Kzz + I_like(Kzz) * self.jitter)
+        Kzz = psd(Dense(Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter))
 
         sqrt = Triangular(sqrt)
         S = sqrt @ sqrt.T
@@ -226,7 +224,8 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         kernel = self.posterior.prior.kernel
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
 
@@ -238,10 +237,10 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
         mut = mean_function(t)
 
         # Lz⁻¹ Kzt
-        Lz_inv_Kzt = solve(Lz, Kzt, Cholesky())
+        Lz_inv_Kzt = solve(Lz, Kzt)
 
         # Kzz⁻¹ Kzt
-        Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt, Cholesky())
+        Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt)
 
         # Ktz Kzz⁻¹ sqrt
         Ktz_Kzz_inv_sqrt = jnp.matmul(Kzz_inv_Kzt.T, sqrt)
@@ -255,7 +254,7 @@ 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 += I_like(covariance) * self.jitter
+        covariance += jnp.eye(covariance.shape[0]) * self.jitter
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -330,7 +329,8 @@ class WhitenedVariationalGaussian(VariationalGaussian[L]):
         kernel = self.posterior.prior.kernel
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
 
         # Unpack test inputs
@@ -341,7 +341,7 @@ class WhitenedVariationalGaussian(VariationalGaussian[L]):
         mut = mean_function(t)
 
         # Lz⁻¹ Kzt
-        Lz_inv_Kzt = solve(Lz, Kzt, Cholesky())
+        Lz_inv_Kzt = solve(Lz, Kzt)
 
         # Ktz Lz⁻ᵀ sqrt
         Ktz_Lz_invT_sqrt = jnp.matmul(Lz_inv_Kzt.T, sqrt)
@@ -355,7 +355,7 @@ 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 += I_like(covariance) * self.jitter
+        covariance += jnp.eye(covariance.shape[0]) * self.jitter
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -441,7 +441,8 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
 
         qu = GaussianDistribution(loc=jnp.atleast_1d(mu.squeeze()), scale=S)
         pu = GaussianDistribution(loc=jnp.atleast_1d(muz.squeeze()), scale=Kzz)
@@ -492,7 +493,8 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
         mu = jnp.matmul(S, natural_vector)
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
 
@@ -501,10 +503,10 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
         mut = mean_function(test_inputs)
 
         # Lz⁻¹ Kzt
-        Lz_inv_Kzt = solve(Lz, Kzt, Cholesky())
+        Lz_inv_Kzt = solve(Lz, Kzt)
 
         # Kzz⁻¹ Kzt
-        Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt, Cholesky())
+        Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt)
 
         # Ktz Kzz⁻¹ L
         Ktz_Kzz_inv_L = jnp.matmul(Kzz_inv_Kzt.T, sqrt)
@@ -518,7 +520,7 @@ 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 += I_like(covariance) * self.jitter
+        covariance += jnp.eye(covariance.shape[0]) * self.jitter
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -592,12 +594,14 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
 
         # S = η₂ - η₁ η₁ᵀ
         S = expectation_matrix - jnp.outer(mu, mu)
-        S = PSD(Dense(S))
-        S += I_like(S) * self.jitter
+        S = psd(Dense(S))
+        S_dense = S.to_dense() + jnp.eye(S.shape[0]) * self.jitter
+        S = psd(Dense(S_dense))
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
 
         qu = GaussianDistribution(loc=jnp.atleast_1d(mu.squeeze()), scale=S)
         pu = GaussianDistribution(loc=jnp.atleast_1d(muz.squeeze()), scale=Kzz)
@@ -636,14 +640,15 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
 
         # S = η₂ - η₁ η₁ᵀ
         S = expectation_matrix - jnp.matmul(mu, mu.T)
-        S = Dense(S) + I_like(S) * self.jitter
-        S = PSD(S)
+        S = Dense(S + jnp.eye(S.shape[0]) * self.jitter)
+        S = psd(S)
 
         # S = sqrt sqrtᵀ
         sqrt = lower_cholesky(S)
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
 
@@ -655,10 +660,10 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
         mut = mean_function(t)
 
         # Lz⁻¹ Kzt
-        Lz_inv_Kzt = solve(Lz, Kzt, Cholesky())
+        Lz_inv_Kzt = solve(Lz, Kzt)
 
         # Kzz⁻¹ Kzt
-        Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt, Cholesky())
+        Kzz_inv_Kzt = solve(Lz.T, Lz_inv_Kzt)
 
         # Ktz Kzz⁻¹ sqrt
         Ktz_Kzz_inv_sqrt = Kzz_inv_Kzt.T @ sqrt
@@ -672,7 +677,7 @@ 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 += I_like(covariance) * self.jitter
+        covariance += jnp.eye(covariance.shape[0]) * self.jitter
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance
@@ -729,13 +734,14 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]):
 
         Kzx = kernel.cross_covariance(z, x)
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = Kzz.to_dense() + jnp.eye(Kzz.shape[0]) * self.jitter
+        Kzz = psd(Dense(Kzz_dense))
 
         # Lz Lzᵀ = Kzz
         Lz = lower_cholesky(Kzz)
 
         # Lz⁻¹ Kzx
-        Lz_inv_Kzx = solve(Lz, Kzx, Cholesky())
+        Lz_inv_Kzx = solve(Lz, Kzx)
 
         # A = Lz⁻¹ Kzt / o
         A = Lz_inv_Kzx / self.posterior.likelihood.obs_stddev.value
@@ -753,14 +759,14 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]):
         Lz_inv_Kzx_diff = jsp.linalg.cho_solve((L, True), jnp.matmul(Lz_inv_Kzx, diff))
 
         # Kzz⁻¹ Kzx (y - μx)
-        Kzz_inv_Kzx_diff = solve(Lz.T, Lz_inv_Kzx_diff, Cholesky())
+        Kzz_inv_Kzx_diff = solve(Lz.T, Lz_inv_Kzx_diff)
 
         Ktt = kernel.gram(t)
         Kzt = kernel.cross_covariance(z, t)
         mut = mean_function(t)
 
         # Lz⁻¹ Kzt
-        Lz_inv_Kzt = solve(Lz, Kzt, Cholesky())
+        Lz_inv_Kzt = solve(Lz, Kzt)
 
         # L⁻¹ Lz⁻¹ Kzt
         L_inv_Lz_inv_Kzt = jsp.linalg.solve_triangular(L, Lz_inv_Kzt, lower=True)
@@ -774,7 +780,7 @@ 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 += I_like(covariance) * self.jitter
+        covariance += jnp.eye(covariance.shape[0]) * self.jitter
 
         return GaussianDistribution(
             loc=jnp.atleast_1d(mean.squeeze()), scale=covariance

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.11.2
+Version: 0.12.0
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -14,11 +14,11 @@ Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
-Requires-Python: <3.13,>=3.10
+Requires-Python: <=3.13,>=3.10
 Requires-Dist: beartype>0.16.1
-Requires-Dist: cola-ml>=0.0.7
 Requires-Dist: flax>=0.10.0
 Requires-Dist: jax>=0.5.0
 Requires-Dist: jaxlib>=0.5.0
@@ -27,6 +27,47 @@ Requires-Dist: numpy>=2.0.0
 Requires-Dist: numpyro
 Requires-Dist: optax>0.2.1
 Requires-Dist: tqdm>4.66.2
+Provides-Extra: dev
+Requires-Dist: absolufy-imports>=0.3.1; extra == 'dev'
+Requires-Dist: autoflake; extra == 'dev'
+Requires-Dist: black; extra == 'dev'
+Requires-Dist: codespell>=2.2.4; extra == 'dev'
+Requires-Dist: coverage>=7.2.2; extra == 'dev'
+Requires-Dist: interrogate>=1.5.0; extra == 'dev'
+Requires-Dist: isort; extra == 'dev'
+Requires-Dist: jupytext; extra == 'dev'
+Requires-Dist: mktestdocs>=0.2.1; extra == 'dev'
+Requires-Dist: networkx; extra == 'dev'
+Requires-Dist: pre-commit>=3.2.2; extra == 'dev'
+Requires-Dist: pytest-beartype; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest-pretty>=1.1.1; extra == 'dev'
+Requires-Dist: pytest-xdist>=3.2.1; extra == 'dev'
+Requires-Dist: pytest>=7.2.2; extra == 'dev'
+Requires-Dist: ruff>=0.6; extra == 'dev'
+Requires-Dist: xdoctest>=1.1.1; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: blackjax>=0.9.6; extra == 'docs'
+Requires-Dist: ipykernel>=6.22.0; extra == 'docs'
+Requires-Dist: ipython>=8.11.0; extra == 'docs'
+Requires-Dist: ipywidgets>=8.0.5; extra == 'docs'
+Requires-Dist: jupytext>=1.14.5; extra == 'docs'
+Requires-Dist: markdown-katex>=202406.1035; extra == 'docs'
+Requires-Dist: matplotlib>=3.7.1; extra == 'docs'
+Requires-Dist: mkdocs-gen-files>=0.5.0; extra == 'docs'
+Requires-Dist: mkdocs-git-authors-plugin>=0.7.0; extra == 'docs'
+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: nbconvert>=7.16.2; extra == 'docs'
+Requires-Dist: networkx>=3.0; extra == 'docs'
+Requires-Dist: pandas>=1.5.3; extra == 'docs'
+Requires-Dist: pymdown-extensions>=10.7.1; extra == 'docs'
+Requires-Dist: scikit-learn>=1.5.1; extra == 'docs'
+Requires-Dist: seaborn>=0.12.2; extra == 'docs'
+Requires-Dist: watermark>=2.3.1; extra == 'docs'
 Description-Content-Type: text/markdown
 
 <!-- <h1 align='center'>GPJax</h1>
@@ -41,12 +82,12 @@ Description-Content-Type: text/markdown
 [![PyPI version](https://badge.fury.io/py/GPJax.svg)](https://badge.fury.io/py/GPJax)
 [![DOI](https://joss.theoj.org/papers/10.21105/joss.04455/status.svg)](https://doi.org/10.21105/joss.04455)
 [![Downloads](https://pepy.tech/badge/gpjax)](https://pepy.tech/project/gpjax)
-[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
+[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
 
 [**Quickstart**](#simple-example)
 | [**Install guide**](#installation)
 | [**Documentation**](https://docs.jaxgaussianprocesses.com/)
-| [**Slack Community**](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
+| [**Slack Community**](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
 
 GPJax aims to provide a low-level interface to Gaussian process (GP) models in
 [Jax](https://github.com/google/jax), structured to give researchers maximum
@@ -81,7 +122,7 @@ behaviours through [this form](https://jaxgaussianprocesses.com/contact/) or rea
 one of the project's [_gardeners_](https://docs.jaxgaussianprocesses.com/GOVERNANCE/#roles).
 
 Feel free to join our [Slack
-Channel](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw),
+Channel](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA),
 where we can discuss the development of GPJax and broader support for Gaussian
 process modelling.
 
@@ -177,14 +218,13 @@ configuration in development mode.
 ```bash
 git clone https://github.com/JaxGaussianProcesses/GPJax.git
 cd GPJax
-hatch env create
-hatch shell
+uv sync --extra dev
 ```
 
 > We recommend you check your installation passes the supplied unit tests:
 >
 > ```python
-> hatch run dev:test
+> uv run pytest --beartype-packages='gpjax'
 > ```
 
 # Citing GPJax

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

@@ -1,29 +1,28 @@
-gpjax/__init__.py,sha256=ylCFMtXwcMS2zxm4pI3KsnRdnX6bdh26TSdTfUh9l9o,1686
+gpjax/__init__.py,sha256=FSrKDFSQ7xDqwQGBWwEPqqjvYxEbhUPPestKLoAPjWA,1686
 gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
-gpjax/distributions.py,sha256=8LWmfmRVHOX29Uy8PkKFi2UhcCiunuu-4TMI_5-krHc,9299
+gpjax/distributions.py,sha256=iKmeQ_NN2CIjRiuOeJlwEGASzGROi4ZCerVi1uY7zRM,7758
 gpjax/fit.py,sha256=R4TIPvBNHYSg9vBVp6is_QYENldRLIU_FklGE85C-aA,15046
-gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
+gpjax/gps.py,sha256=-Log0pcU8qmB5fUxfzoNjD0S64gpiypAjFzjGXX6w7I,30301
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
 gpjax/likelihoods.py,sha256=99oTZoWld1M7vxgGM0pNY5Hnt2Ajd2lQNqawzrLmwtk,9308
-gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
 gpjax/mean_functions.py,sha256=-sVYO1_LWE8f34rllUOuaT5sgGGAdxo99v5kRo2d4oM,6490
 gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
-gpjax/objectives.py,sha256=I_ZqnwTNYIAUAZ9KQNenIl0ish1jDOXb7KaNmjz3Su4,15340
-gpjax/parameters.py,sha256=H-DiXmotdBZCbf-GOjRaJoS_isk3GgFrpKHTq5GpnoA,6998
+gpjax/objectives.py,sha256=Tm36h8fz_nWkZPlufMQzZWKK1ytrtT9yvvP8YdxYKNw,15359
+gpjax/parameters.py,sha256=qIEqyMKNd2n2Ak15PisCmqhX5qhsoRgng_s4doL96rE,7044
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=Y9J1H91tXPm_hMy3ri_PgjAxqc_3r-BqKV83HRvB_m4,28295
+gpjax/variational_families.py,sha256=rE3LarwIAkvDvLlWrz8Ww6BUBz88YHdV4ceY97r3IBw,28637
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
-gpjax/kernels/base.py,sha256=wXsrpm5ofy9S5MNgUkJk4lx2umcIJL6dDNhXY7cmTGk,11616
+gpjax/kernels/base.py,sha256=hOUXwarspDFnuI2_QreyIVPdz2fzRVJj4p3Zdu1touw,11606
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
 gpjax/kernels/approximations/rff.py,sha256=VbitjNuahFE5_IvCj1A0SxHhJXU0O0Qq0FMMVq8xA3E,4125
 gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
-gpjax/kernels/computations/base.py,sha256=zzabLN_-FkTWo6cBYjzjvUGYa7vrYyHxyrhQZxLgHBk,3651
-gpjax/kernels/computations/basis_functions.py,sha256=zY4rUDZDLOYvQPY9xosRmCLPdiXfbzAN5GICjQhFOis,2528
-gpjax/kernels/computations/constant_diagonal.py,sha256=_4fIFPq76Z416-9dIr8N335WP291dGluO-RqqUsK9ZY,2058
+gpjax/kernels/computations/base.py,sha256=L6K0roxZbrYeJKxEw-yaTiK9Mtcv0YtZfWI2Xnau7i8,3616
+gpjax/kernels/computations/basis_functions.py,sha256=MPSo40NEx_ngnSLTa9ntVJzma_jugvm5dMpZd5MtG5M,2490
+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=z2JpUue7oY-pL-c0Pc6Bngv_IJCR6z4MaW7kN0wgjxk,1803
+gpjax/kernels/computations/diagonal.py,sha256=k1KqW0DwWRIBvbb7jzcKktXRfhXbcos3ncWrFplJ4W0,1768
 gpjax/kernels/computations/eigen.py,sha256=w7I7LK42j0ouchHCI1ltXx0lpwqvK1bRb4HclnF3rKs,1936
 gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3J9_fWz1NlY,790
 gpjax/kernels/non_euclidean/graph.py,sha256=K4WIdX-dx1SsWuNHZnNjHFw8ElKZxGcReUiA3w4aCOI,4204
@@ -43,7 +42,11 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 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-0.11.2.dist-info/METADATA,sha256=lTQVlrUbkxI7fU9Gdnac_eoNRyjCHEoEuEvvWbKmDqM,8558
-gpjax-0.11.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.11.2.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.11.2.dist-info/RECORD,,
+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/lower_cholesky.py

@@ -1,69 +0,0 @@
-# Copyright 2023 The GPJax Contributors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from cola.annotations import PSD
-from cola.fns import dispatch
-from cola.ops.operator_base import LinearOperator
-from cola.ops.operators import (
-    BlockDiag,
-    Diagonal,
-    Identity,
-    Kronecker,
-    Triangular,
-)
-import jax.numpy as jnp
-
-# TODO: Once this functionality is supported in CoLA, remove this.
-
-
-@dispatch
-def lower_cholesky(A: LinearOperator) -> Triangular:  # noqa: F811
-    """Returns the lower Cholesky factor of a linear operator.
-
-    Args:
-        A: The input linear operator.
-
-    Returns:
-        Triangular: The lower Cholesky factor of A.
-    """
-
-    if PSD not in A.annotations:
-        raise ValueError(
-            "Expected LinearOperator to be PSD, did you forget to use cola.PSD?"
-        )
-
-    return Triangular(jnp.linalg.cholesky(A.to_dense()), lower=True)
-
-
-@lower_cholesky.dispatch
-def _(A: Diagonal):  # noqa: F811
-    return Diagonal(jnp.sqrt(A.diag))
-
-
-@lower_cholesky.dispatch
-def _(A: Identity):  # noqa: F811
-    return A
-
-
-@lower_cholesky.dispatch
-def _(A: Kronecker):  # noqa: F811
-    return Kronecker(*[lower_cholesky(Ai) for Ai in A.Ms])
-
-
-@lower_cholesky.dispatch
-def _(A: BlockDiag):  # noqa: F811
-    return BlockDiag(
-        *[lower_cholesky(Ai) for Ai in A.Ms], multiplicities=A.multiplicities
-    )