gpjax 0.9.1__py3-none-any.whl → 0.9.2__py3-none-any.whl

gpjax/__init__.py

@@ -40,7 +40,7 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.9.1"
+__version__ = "0.9.2"
 
 __all__ = [
     "base",

gpjax/gps.py

@@ -17,6 +17,7 @@ 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
@@ -530,6 +531,7 @@ 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.
 
@@ -563,6 +565,11 @@ 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
@@ -588,7 +595,7 @@ class ConjugatePosterior(AbstractPosterior[P, GL]):
         canonical_weights = solve(
             Sigma,
             y + eps - jnp.inner(Phi, fourier_weights),
-            Cholesky(),
+            solver_algorithm,
         )  #  [N, B]
 
         def sample_fn(test_inputs: Float[Array, "n D"]) -> Float[Array, "n B"]:

{gpjax-0.9.1.dist-info → gpjax-0.9.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: gpjax
-Version: 0.9.1
+Version: 0.9.2
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues

{gpjax-0.9.1.dist-info → gpjax-0.9.2.dist-info}/RECORD

@@ -1,9 +1,9 @@
-gpjax/__init__.py,sha256=WI2T3AUoqBPcFQe8_pflRkV1k516x6ljUJsqD_7FUBY,1697
+gpjax/__init__.py,sha256=Bx5JFaveeVk3qJMTzbmrKOFy0U7fNcQ_JnVo5m0ACGA,1697
 gpjax/citation.py,sha256=R4Pmvjt0ndA0avEDSvIbxDxKapkRRYXWX7RRWBvZCRQ,5306
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
 gpjax/distributions.py,sha256=zxkSEZIlTg0PHvvgj0BQuIFEg-ugx6_NkEwSsbqWUM0,9325
 gpjax/fit.py,sha256=OHv8jUHxa1ndpqMERSDRtYtUDzubk9rMPVIhfCiIH5Q,11551
-gpjax/gps.py,sha256=6v0P6fw9CmyHfE8S3WwxCE2qECQ_kaq_SQWHkz6UIPE,30612
+gpjax/gps.py,sha256=NO18geRfcjo4mA3PGkuGont_Mj_yRqfvWzJqYmoKwiY,31225
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
 gpjax/likelihoods.py,sha256=Uh4kgLTod8ODw178L--G3w4olpm9XvCdcAZ8l7FwkF4,9255
 gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
@@ -56,7 +56,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 gpjax/kernels/stationary/rbf.py,sha256=G13gg5phO7ite7D9QgoCy7gB2_y0FM6GZhgFW4RL6Xw,1734
 gpjax/kernels/stationary/utils.py,sha256=Xa9EEnxgFqEi08ZSFAZYYHhJ85_3Ac-ZUyUk18B63M4,2225
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-gpjax-0.9.1.dist-info/METADATA,sha256=aR6Xzmn8HdUzxvE7cdKuL6lMI21kHoDQaaUbM9wkzFw,9976
-gpjax-0.9.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-gpjax-0.9.1.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax-0.9.1.dist-info/RECORD,,
+gpjax-0.9.2.dist-info/METADATA,sha256=JWT3cDW7onuKnTYUGqa15WxG4L7oEboJKPHYyAggYZ0,9976
+gpjax-0.9.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+gpjax-0.9.2.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
+gpjax-0.9.2.dist-info/RECORD,,