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

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,12 +32,19 @@ 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.linalg.utils import add_jitter
 from gpjax.mean_functions import AbstractMeanFunction
 from gpjax.parameters import (
     LowerTriangular,
     Real,
-    Static,
 )
 from gpjax.typing import (
     Array,
@@ -112,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
@@ -179,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(Kzz + I_like(Kzz) * 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.
@@ -217,37 +220,38 @@ 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 += I_like(Kzz) * 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, 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)
+        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 = (
@@ -255,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 += I_like(covariance) * 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
@@ -330,7 +337,8 @@ class WhitenedVariationalGaussian(VariationalGaussian[L]):
         kernel = self.posterior.prior.kernel
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
 
         # Unpack test inputs
@@ -341,7 +349,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 +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 += I_like(covariance) * 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,8 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * 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)
         pu = GaussianDistribution(loc=jnp.atleast_1d(muz.squeeze()), scale=Kzz)
@@ -475,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(
@@ -492,7 +504,8 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
         mu = jnp.matmul(S, natural_vector)
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
 
@@ -501,10 +514,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 +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 += I_like(covariance) * 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
@@ -554,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.
@@ -592,12 +608,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 = add_jitter(S.to_dense(), self.jitter)
+        S = psd(Dense(S_dense))
 
         muz = mean_function(z)
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * 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)
         pu = GaussianDistribution(loc=jnp.atleast_1d(muz.squeeze()), scale=Kzz)
@@ -636,14 +654,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(add_jitter(S, self.jitter))
+        S = psd(S)
 
         # S = sqrt sqrtᵀ
         sqrt = lower_cholesky(S)
 
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), self.jitter)
+        Kzz = psd(Dense(Kzz_dense))
         Lz = lower_cholesky(Kzz)
         muz = mean_function(z)
 
@@ -655,10 +674,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 +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 += I_like(covariance) * 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
@@ -729,13 +751,14 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]):
 
         Kzx = kernel.cross_covariance(z, x)
         Kzz = kernel.gram(z)
-        Kzz += I_like(Kzz) * self.jitter
+        Kzz_dense = add_jitter(Kzz.to_dense(), 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 +776,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 +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 += I_like(covariance) * 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.11.2.dist-info → gpjax-0.12.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.11.2
+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
@@ -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.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.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'
+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,22 +122,13 @@ 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.
 
-
-## 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
 

gpjax-0.12.2.dist-info/RECORD

@@ -0,0 +1,52 @@
+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=I2sJVuKZii_d7MEcelHIivfM8ExYGMgdBuKKOT7Dw-A,15326
+gpjax/gps.py,sha256=ipaeYMnPffhKK_JsEHe4fF8GmolQIjXB1YbyfUIL8H4,30118
+gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
+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=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=TJGGkwkE805X4PQb-C32FxvD9B_OsFLWf6I-ZZvOUWk,29628
+gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
+gpjax/kernels/base.py,sha256=4Lx8y3kPX4WqQZGRGAsBkqn_i6FlfoAhSn9Tv415xuQ,11551
+gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
+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=_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=NTHm-cn-RepYuXFrvXo2ih7Gtu1YR_pAg4Jb7IhE_o8,1930
+gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3J9_fWz1NlY,790
+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=cqb8sqaNwW3fEbrA7MY9OF2KJFTkxHhqwmQtABE3G8w,5408
+gpjax/kernels/nonstationary/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
+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=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=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,,

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

gpjax-0.11.2.dist-info/RECORD

@@ -1,49 +0,0 @@
-gpjax/__init__.py,sha256=ylCFMtXwcMS2zxm4pI3KsnRdnX6bdh26TSdTfUh9l9o,1686
-gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
-gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
-gpjax/distributions.py,sha256=8LWmfmRVHOX29Uy8PkKFi2UhcCiunuu-4TMI_5-krHc,9299
-gpjax/fit.py,sha256=R4TIPvBNHYSg9vBVp6is_QYENldRLIU_FklGE85C-aA,15046
-gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
-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/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/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
-gpjax/kernels/base.py,sha256=wXsrpm5ofy9S5MNgUkJk4lx2umcIJL6dDNhXY7cmTGk,11616
-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/dense.py,sha256=vnW6XKQe4_gzpXRWTctxhgMA9-9TebdtiXzAqh_-j6g,1392
-gpjax/kernels/computations/diagonal.py,sha256=z2JpUue7oY-pL-c0Pc6Bngv_IJCR6z4MaW7kN0wgjxk,1803
-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
-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/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
-gpjax/kernels/nonstationary/polynomial.py,sha256=arP8DK0jnBOaayDWcFvHF0pdu9FVhwzXdqjnHUAL2VI,3293
-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/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,,