Merge pull request #454 from JaxGaussianProcesses/avoid_cg

Quick fix to stop automatic switch to CG
JaxGaussianProcesses · Jun 18, 2024 · 5734eac · 5734eac
2 parents ad4fdac + 9049189
commit 5734eac
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gpjax/gps.py b/gpjax/gps.py
@@ -532,7 +532,7 @@ def predict(
         mean_t = self.prior.mean_function(t)
         Ktt = self.prior.kernel.gram(t)
         Kxt = self.prior.kernel.cross_covariance(x, t)
-        Sigma_inv_Kxt = cola.solve(Sigma, Kxt)
+        Sigma_inv_Kxt = cola.solve(Sigma, Kxt, Cholesky())
 
         # μt  +  Ktx (Kxx + Io²)⁻¹ (y  -  μx)
         mean = mean_t + jnp.matmul(Sigma_inv_Kxt.T, y - mx)