Test Gaussian cov is PSD, bump workflows to python 3.10

.github/workflows/benchmarks.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.8"]
+        python-version: ["3.10"]
     steps:
       - name: Checkout the branch
         uses: actions/checkout@v3.5.2
@@ -24,7 +24,7 @@ jobs:
       - name: Setup Python
         uses: actions/setup-python@v3
         with:
-          python-version: "3.8"
+          python-version: "3.10"
 
       - name: Install dependencies
         run: |

.github/workflows/build_docs.yml

@@ -19,7 +19,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.8"]
+        python-version: ["3.10"]
 
     steps:
       # Grap the latest commit from the branch

.github/workflows/integration.yml

@@ -13,7 +13,7 @@ jobs:
       matrix:
         # Select the Python versions to test against
         os: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.10", "3.11"]
       fail-fast: true
     steps:
       - name: Check out the code

.github/workflows/test_docs.yml

@@ -17,7 +17,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.8"]
+        python-version: ["3.10"]
 
     steps:
       # Grap the latest commit from the branch

gpjax/lower_cholesky.py

@@ -16,7 +16,7 @@
 import cola
 import jax.numpy as jnp
 
-# TODO: Add lower_cholesky for other linear operators.
+# TODO: Once this functionality is supported in CoLA, remove this.
 
 
 @cola.dispatch

tests/test_gaussian_distribution.py

@@ -60,6 +60,9 @@ def test_array_arguments(n: int) -> None:
     assert approx_equal(dist.stddev(), jnp.sqrt(covariance.diagonal()))
     assert approx_equal(dist.covariance(), covariance)
 
+    assert isinstance(dist.scale, Dense)
+    assert cola.PSD in dist.scale.annotations
+
     y = jr.uniform(_key, shape=(n,))
 
     tfp_dist = MultivariateNormalFullCovariance(loc=mean, covariance_matrix=covariance)
@@ -74,9 +77,14 @@ def test_diag_linear_operator(n: int) -> None:
     mean = jr.uniform(key_mean, shape=(n,))
     diag = jr.uniform(key_diag, shape=(n,))
 
+    # We purosely forget to add a PSD annotation to the diagonal matrix.
     dist_diag = GaussianDistribution(loc=mean, scale=Diagonal(diag**2))
     tfp_dist = MultivariateNormalDiag(loc=mean, scale_diag=diag)
 
+    # We check that the PSD annotation is added automatically.
+    assert isinstance(dist_diag.scale, Diagonal)
+    assert cola.PSD in dist_diag.scale.annotations
+
     assert approx_equal(dist_diag.mean(), tfp_dist.mean())
     assert approx_equal(dist_diag.mode(), tfp_dist.mode())
     assert approx_equal(dist_diag.entropy(), tfp_dist.entropy())

tests/test_kernels/test_computation.py

@@ -9,6 +9,7 @@
 from gpjax.kernels.computations import (
     ConstantDiagonalKernelComputation,
     DiagonalKernelComputation,
+    BasisFunctionComputation,
 )
 from gpjax.kernels.nonstationary import (
     Linear,
@@ -77,4 +78,4 @@ def test_change_computation(kernel):
     assert jnp.allclose(constant_entries, constant_entries[0])
 
     # All the off diagonal entries should be zero
-    assert jnp.allclose(constant_diagonal_matrix - jnp.diag(constant_entries), 0.0)
+    assert jnp.allclose(constant_diagonal_matrix - jnp.diag(constant_entries), 0.0)