gpjax 0.9.0__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.0"
+__version__ = "0.9.2"
 
 __all__ = [
     "base",

gpjax/decision_making/test_functions/continuous_functions.py

@@ -139,12 +139,7 @@ class LogarithmicGoldsteinPrice(AbstractContinuousTestFunction):
         x1 = 4.0 * x[:, 0] - 2.0
         x2 = 4.0 * x[:, 1] - 2.0
         a = 1.0 + (x1 + x2 + 1.0) ** 2 * (
-            19.0
-            - 14.0 * x1
-            + 3.0 * (x1**2)
-            - 14.0 * x2
-            + 6.0 * x1 * x2
-            + 3.0 * (x2**2)
+            19.0 - 14.0 * x1 + 3.0 * (x1**2) - 14.0 * x2 + 6.0 * x1 * x2 + 3.0 * (x2**2)
         )
         b = 30.0 + (2.0 * x1 - 3.0 * x2) ** 2 * (
             18.0

gpjax/distributions.py

@@ -162,9 +162,7 @@ class GaussianDistribution(tfd.Distribution):
 
         return vmap(affine_transformation)(Z)
 
-    def sample(
-        self, seed: KeyArray, sample_shape: Tuple[int, ...]
-    ):  # pylint: disable=useless-super-delegation
+    def sample(self, seed: KeyArray, sample_shape: Tuple[int, ...]):  # pylint: disable=useless-super-delegation
         r"""See `Distribution.sample`."""
         return self._sample_n(
             seed, sample_shape[0]

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
@@ -146,20 +147,17 @@ class Prior(AbstractPrior[M, K]):
     if tp.TYPE_CHECKING:
 
         @tp.overload
-        def __mul__(self, other: GL) -> "ConjugatePosterior[Prior[M, K], GL]":
-            ...
+        def __mul__(self, other: GL) -> "ConjugatePosterior[Prior[M, K], GL]": ...
 
         @tp.overload
         def __mul__(  # noqa: F811
             self, other: NGL
-        ) -> "NonConjugatePosterior[Prior[M, K], NGL]":
-            ...
+        ) -> "NonConjugatePosterior[Prior[M, K], NGL]": ...
 
         @tp.overload
         def __mul__(  # noqa: F811
             self, other: L
-        ) -> "AbstractPosterior[Prior[M, K], L]":
-            ...
+        ) -> "AbstractPosterior[Prior[M, K], L]": ...
 
     def __mul__(self, other):  # noqa: F811
         r"""Combine the prior with a likelihood to form a posterior distribution.
@@ -195,20 +193,17 @@ class Prior(AbstractPrior[M, K]):
     if tp.TYPE_CHECKING:
 
         @tp.overload
-        def __rmul__(self, other: GL) -> "ConjugatePosterior[Prior[M, K], GL]":
-            ...
+        def __rmul__(self, other: GL) -> "ConjugatePosterior[Prior[M, K], GL]": ...
 
         @tp.overload
         def __rmul__(  # noqa: F811
             self, other: NGL
-        ) -> "NonConjugatePosterior[Prior[M, K], NGL]":
-            ...
+        ) -> "NonConjugatePosterior[Prior[M, K], NGL]": ...
 
         @tp.overload
         def __rmul__(  # noqa: F811
             self, other: L
-        ) -> "AbstractPosterior[Prior[M, K], L]":
-            ...
+        ) -> "AbstractPosterior[Prior[M, K], L]": ...
 
     def __rmul__(self, other):  # noqa: F811
         r"""Combine the prior with a likelihood to form a posterior distribution.
@@ -536,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.
 
@@ -569,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
@@ -594,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"]:
@@ -722,15 +723,13 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
 
 
 @tp.overload
-def construct_posterior(prior: P, likelihood: GL) -> ConjugatePosterior[P, GL]:
-    ...
+def construct_posterior(prior: P, likelihood: GL) -> ConjugatePosterior[P, GL]: ...
 
 
 @tp.overload
 def construct_posterior(  # noqa: F811
     prior: P, likelihood: NGL
-) -> NonConjugatePosterior[P, NGL]:
-    ...
+) -> NonConjugatePosterior[P, NGL]: ...
 
 
 def construct_posterior(prior, likelihood):  # noqa: F811

gpjax/kernels/approximations/rff.py

@@ -1,4 +1,5 @@
-"""Compute Random Fourier Feature (RFF) kernel approximations.  """
+"""Compute Random Fourier Feature (RFF) kernel approximations."""
+
 import beartype.typing as tp
 import jax.random as jr
 from jaxtyping import Float

gpjax/kernels/computations/base.py

@@ -74,8 +74,7 @@ class AbstractKernelComputation:
     @abc.abstractmethod
     def _cross_covariance(
         self, kernel: K, x: Num[Array, "N D"], y: Num[Array, "M D"]
-    ) -> Float[Array, "N M"]:
-        ...
+    ) -> Float[Array, "N M"]: ...
 
     def cross_covariance(
         self, kernel: K, x: Num[Array, "N D"], y: Num[Array, "M D"]

gpjax/kernels/computations/eigen.py

@@ -26,7 +26,8 @@ from gpjax.kernels.computations.base import AbstractKernelComputation
 from gpjax.typing import Array
 
 Kernel = tp.TypeVar(
-    "Kernel", bound="gpjax.kernels.non_euclidean.graph.GraphKernel"  # noqa: F821
+    "Kernel",
+    bound="gpjax.kernels.non_euclidean.graph.GraphKernel",  # noqa: F821
 )
 
 

gpjax/kernels/stationary/matern12.py

@@ -38,6 +38,7 @@ class Matern12(StationaryKernel):
     k(x, y) = \sigma^2\exp\Bigg(-\frac{\lvert x-y \rvert}{2\ell^2}\Bigg)
     $$
     """
+
     name: str = "Matérn12"
 
     def __call__(self, x: Float[Array, " D"], y: Float[Array, " D"]) -> ScalarFloat:

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

@@ -1,36 +1,39 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: gpjax
-Version: 0.9.0
+Version: 0.9.2
 Summary: Gaussian processes in JAX.
-Home-page: https://github.com/JaxGaussianProcesses/GPJax
-License: Apache-2.0
+Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
+Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
+Project-URL: Source, https://github.com/JaxGaussianProcesses/GPJax
+Author-email: Thomas Pinder <tompinder@live.co.uk>
+License-Expression: Apache-2.0
+License-File: LICENSE
 Keywords: gaussian-processes jax machine-learning bayesian
-Author: Thomas Pinder
-Author-email: tompinder@live.co.uk
-Requires-Python: >=3.10,<3.12
-Classifier: License :: OSI Approved :: Apache Software License
-Classifier: Programming Language :: Python :: 3
+Classifier: Development Status :: 4 - Beta
+Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: beartype (>=0.16.1,<0.17.0)
-Requires-Dist: cola-ml (==0.0.5)
-Requires-Dist: flax (>=0.8.4,<0.9.0)
-Requires-Dist: jax (<0.4.28)
-Requires-Dist: jaxlib (<0.4.28)
-Requires-Dist: jaxopt (>=0.8.3,<0.9.0)
-Requires-Dist: jaxtyping (>=0.2.10,<0.3.0)
-Requires-Dist: numpy (<2.0.0)
-Requires-Dist: optax (>=0.2.1,<0.3.0)
-Requires-Dist: tensorflow-probability (>=0.24.0,<0.25.0)
-Requires-Dist: tqdm (>=4.66.2,<5.0.0)
-Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
-Project-URL: Repository, https://github.com/JaxGaussianProcesses/GPJax
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Requires-Python: <3.13,>=3.10
+Requires-Dist: beartype>0.16.1
+Requires-Dist: cola-ml==0.0.5
+Requires-Dist: flax>=0.8.5
+Requires-Dist: jax<0.4.28
+Requires-Dist: jaxlib<0.4.28
+Requires-Dist: jaxopt==0.8.2
+Requires-Dist: jaxtyping>0.2.10
+Requires-Dist: numpy<2.0.0
+Requires-Dist: optax>0.2.1
+Requires-Dist: tensorflow-probability>=0.24.0
+Requires-Dist: tqdm>4.66.2
 Description-Content-Type: text/markdown
 
 <!-- <h1 align='center'>GPJax</h1>
 <h2 align='center'>Gaussian processes in Jax.</h2> -->
 <p align="center">
-<img width="700" height="300" src="https://raw.githubusercontent.com/JaxGaussianProcesses/GPJax/main/docs/_static/gpjax_logo.svg" alt="GPJax's logo">
+<img width="700" height="300" src="https://raw.githubusercontent.com/JaxGaussianProcesses/GPJax/main/docs/static/gpjax_logo.svg" alt="GPJax's logo">
 </p>
 
 [![codecov](https://codecov.io/gh/JaxGaussianProcesses/GPJax/branch/master/graph/badge.svg?token=DM1DRDASU2)](https://codecov.io/gh/JaxGaussianProcesses/GPJax)
@@ -234,13 +237,14 @@ configuration in development mode.
 ```bash
 git clone https://github.com/JaxGaussianProcesses/GPJax.git
 cd GPJax
-poetry install
+hatch env create
+hatch shell
 ```
 
 > We recommend you check your installation passes the supplied unit tests:
 >
 > ```python
-> poetry run pytest
+> hatch run dev:test
 > ```
 
 # Citing GPJax
@@ -261,4 +265,3 @@ If you use GPJax in your research, please cite our [JOSS paper](https://joss.the
   journal = {Journal of Open Source Software}
 }
 ```
-

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

@@ -1,36 +1,43 @@
-LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax/__init__.py,sha256=ZCiDPPFtUkWm4x1f8FXnkGcXmFf1d8JZBP-gy0_SLCU,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=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
+gpjax/mean_functions.py,sha256=et2HzlsYJNViBvTohF2wZYgCWQfDX4KboYeO7egMR1c,6420
+gpjax/objectives.py,sha256=XwkPyL_iovTNKpKGVNt0Lt2_OMTJitSPhuyCtUrJpbc,15383
+gpjax/parameters.py,sha256=Z4Wy3gEzPZG23-dtqC437_ZWnd_sPe9LcLCKn21ZBvA,4886
+gpjax/scan.py,sha256=mtMsg8yLdkVuOYeTHLnATPGfGDnCMAQNdUA-FJlpfLs,5475
+gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
+gpjax/variational_families.py,sha256=Eik5CCU7qH7_7cacpZ-1lIXm4tElELwSYfVw-n0rI20,27742
 gpjax/decision_making/__init__.py,sha256=SDuPQl80lJ7nhfRsiB_7c22wCMiQO5ehSNohxUGnB7w,2170
 gpjax/decision_making/decision_maker.py,sha256=S4pOXrWcEHy0NDA0gfWzhk7pG0NJfaPpMXvq03yTy0g,13915
 gpjax/decision_making/posterior_handler.py,sha256=UgXf1Gu7GMh2YDSmiSWJIzmWlFW06KTS44HYz3mazZQ,5905
 gpjax/decision_making/search_space.py,sha256=bXwtMOhHZ2klnABpXm5Raxe7b0NTRDjo_cN3ecbk53Y,3545
+gpjax/decision_making/utility_maximizer.py,sha256=VT2amwSJbB64IL_MiWNl9ZgjcqO757qK6NW2gUBKsqs,5965
+gpjax/decision_making/utils.py,sha256=5j1GO5kcmG2laZR39NjhqgEjRekAWWzrnREv_5Zct_Y,2367
 gpjax/decision_making/test_functions/__init__.py,sha256=GDCY9_kaAnxDWwzo1FkdxnDx-80MErAHchbGybT9xYs,1109
-gpjax/decision_making/test_functions/continuous_functions.py,sha256=BdqqPfr4E--LWM_xKTrwEEOLBfRKqjcvd251RR5brCU,5672
+gpjax/decision_making/test_functions/continuous_functions.py,sha256=oL5ZQkvmbC3u9rEvSYI2DRAN3r7Ynf7wRZQlUWjKjt0,5612
 gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=cfo3xQOzB5ajMjjl0YFfNlJClkAcY7ZbT23UyBYEofQ,2955
 gpjax/decision_making/utility_functions/__init__.py,sha256=xXI-4JKWAfTJ7XZ1vRDpqtb91MNzSPD0lP6xo0tOc7o,1445
 gpjax/decision_making/utility_functions/base.py,sha256=FOqrsRDmtHiCVl6IHr12-AEYBLStzMT5EBs-F92e1Og,3882
 gpjax/decision_making/utility_functions/expected_improvement.py,sha256=H6hjC-lj1oiHf2BomeQqroORQ7vtcOngiDAWxRwkNbg,4481
 gpjax/decision_making/utility_functions/probability_of_improvement.py,sha256=O_rHH1yR34JJlpAueSDJ_yo95fPI2aAGkwphS8snBYk,5220
 gpjax/decision_making/utility_functions/thompson_sampling.py,sha256=S-Yyn-9jsKkaXTvKFBP4sG_eCCKApGbHao5RR5tqXAo,4353
-gpjax/decision_making/utility_maximizer.py,sha256=VT2amwSJbB64IL_MiWNl9ZgjcqO757qK6NW2gUBKsqs,5965
-gpjax/decision_making/utils.py,sha256=5j1GO5kcmG2laZR39NjhqgEjRekAWWzrnREv_5Zct_Y,2367
-gpjax/distributions.py,sha256=X48FJr3reop9maherdMVt7-XZOm2f26T8AJt_IKM_oE,9339
-gpjax/fit.py,sha256=OHv8jUHxa1ndpqMERSDRtYtUDzubk9rMPVIhfCiIH5Q,11551
-gpjax/gps.py,sha256=S5_TRG8ki40zpEsizbZ_S9kEppjGcoTEXzFZxVON0WI,30692
-gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
-gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
-gpjax/kernels/approximations/rff.py,sha256=swo7aoqvmwXff2hIcw8HFxD7kJG0IgfVru5XXzxaeys,4127
 gpjax/kernels/base.py,sha256=abkj3zidsBs7YSkYEfjeJ5jTs1YyDCPoBM2ZzqaqrgI,11561
+gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
+gpjax/kernels/approximations/rff.py,sha256=4kD1uocjHmxkLgvf4DxB4_Gy7iefdPgnWiZB3jDiExI,4126
 gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
-gpjax/kernels/computations/base.py,sha256=6L96GwL_6jkjLdvLv3CSAnqK2wJ-vkeVxBfjq7rjqVc,3659
+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=Kx4f59u1r7srlcBVh7z-kfll_K7Km53caAIZYU-5bQ0,1931
+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
@@ -40,7 +47,7 @@ gpjax/kernels/nonstationary/linear.py,sha256=UKDHFCQzKWDMYo76qcb5-ujjnP2_iL-1tcN
 gpjax/kernels/nonstationary/polynomial.py,sha256=yTGobMPbCnKlj4PiQPSXEkWNrj2sjg_x9zFsnFa_j4E,3257
 gpjax/kernels/stationary/__init__.py,sha256=j4BMTaQlIx2kNAT1Dkf4iO2rm-f7_oSVWNrk1bN0tqE,1406
 gpjax/kernels/stationary/base.py,sha256=pQNkMo-E4bIT4tNfb7JvFJZC6fIIXNErsT1iQopFlAA,7063
-gpjax/kernels/stationary/matern12.py,sha256=kFnCxow-xURxxGyUS15G9pUC91X8WPpyCJWQh-xNe8c,1806
+gpjax/kernels/stationary/matern12.py,sha256=b2vQCUqhd9NJK84L2RYjpI597uxy_7xgwsjS35Gc958,1807
 gpjax/kernels/stationary/matern32.py,sha256=ZVYbUIQhvKpriC7abH8wV6Pk-mRoxtl3e2YYwH-KM5Y,2000
 gpjax/kernels/stationary/matern52.py,sha256=xfMYbY7MXxgMECtA2qVT5I8HoDGzGxygUvduGT3_Gvs,2053
 gpjax/kernels/stationary/periodic.py,sha256=IAbCxURtJEHGdmYzbdrsqRZ3zJ8F8tGQF9O7sggafZk,3598
@@ -49,16 +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/likelihoods.py,sha256=Uh4kgLTod8ODw178L--G3w4olpm9XvCdcAZ8l7FwkF4,9255
-gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
-gpjax/mean_functions.py,sha256=et2HzlsYJNViBvTohF2wZYgCWQfDX4KboYeO7egMR1c,6420
-gpjax/objectives.py,sha256=XwkPyL_iovTNKpKGVNt0Lt2_OMTJitSPhuyCtUrJpbc,15383
-gpjax/parameters.py,sha256=Z4Wy3gEzPZG23-dtqC437_ZWnd_sPe9LcLCKn21ZBvA,4886
-gpjax/progress_bar.py,sha256=ErXoJ-jjEPOediq_clgG2OC0UTptpP0MtVM1a2ZLC_o,4418
-gpjax/scan.py,sha256=mtMsg8yLdkVuOYeTHLnATPGfGDnCMAQNdUA-FJlpfLs,5475
-gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=Eik5CCU7qH7_7cacpZ-1lIXm4tElELwSYfVw-n0rI20,27742
-gpjax-0.9.0.dist-info/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax-0.9.0.dist-info/METADATA,sha256=G7J-o4TCBUL91jmsNFHOzJ6f4pxMByGfExUnC6rcx1E,9844
-gpjax-0.9.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-gpjax-0.9.0.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,,

{gpjax-0.9.0.dist-info → gpjax-0.9.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.0
+Generator: hatchling 1.25.0
 Root-Is-Purelib: true
 Tag: py3-none-any

gpjax/progress_bar.py

@@ -1,131 +0,0 @@
-# Copyright 2023 The JaxGaussianProcesses 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 beartype.typing import (
-    Any,
-    Callable,
-    Union,
-)
-from jax import lax
-from jax.experimental import host_callback
-from jaxtyping import Float  # noqa: TCH002
-from tqdm.auto import tqdm
-
-from gpjax.typing import Array  # noqa: TCH001
-
-
-def progress_bar(num_iters: int, log_rate: int) -> Callable:
-    r"""Progress bar decorator for the body function of a `jax.lax.scan`.
-
-    Example:
-    ```python
-        >>> import jax.numpy as jnp
-        >>> import jax
-        >>>
-        >>> carry = jnp.array(0.0)
-        >>> iteration_numbers = jnp.arange(100)
-        >>>
-        >>> @progress_bar(num_iters=iteration_numbers.shape[0], log_rate=10)
-        >>> def body_func(carry, x):
-        >>>    return carry, x
-        >>>
-        >>> carry, _ = jax.lax.scan(body_func, carry, iteration_numbers)
-    ```
-
-    Adapted from this [excellent blog post](https://www.jeremiecoullon.com/2021/01/29/jax_progress_bar/).
-
-    Might be nice in future to directly create a general purpose `verbose scan`
-    inplace of a for a jax.lax.scan, that takes the same arguments as a jax.lax.scan,
-    but prints a progress bar.
-    """
-    tqdm_bars = {}
-    remainder = num_iters % log_rate
-
-    """Define a tqdm progress bar."""
-    tqdm_bars[0] = tqdm(range(num_iters))
-    tqdm_bars[0].set_description("Compiling...", refresh=True)
-
-    def _update_tqdm(args: Any, transform: Any) -> None:
-        r"""Update the tqdm progress bar with the latest objective value."""
-        value, iter_num = args
-        tqdm_bars[0].set_description("Running", refresh=False)
-        tqdm_bars[0].update(iter_num)
-        tqdm_bars[0].set_postfix({"Value": f"{value: .2f}"})
-
-    def _close_tqdm(args: Any, transform: Any) -> None:
-        r"""Close the tqdm progress bar."""
-        tqdm_bars[0].close()
-
-    def _callback(cond: bool, func: Callable, arg: Any) -> None:
-        r"""Callback a function for a given argument if a condition is true."""
-        dummy_result = 0
-
-        def _do_callback(_) -> int:
-            """Perform the callback."""
-            return host_callback.id_tap(func, arg, result=dummy_result)
-
-        def _not_callback(_) -> int:
-            """Do nothing."""
-            return dummy_result
-
-        _ = lax.cond(cond, _do_callback, _not_callback, operand=None)
-
-    def _update_progress_bar(value: Float[Array, "1"], iter_num: int) -> None:
-        r"""Update the tqdm progress bar."""
-        # Conditions for iteration number
-        is_multiple: bool = (iter_num % log_rate == 0) & (
-            iter_num != num_iters - remainder
-        )
-        is_remainder: bool = iter_num == num_iters - remainder
-        is_last: bool = iter_num == num_iters - 1
-
-        # Update progress bar, if multiple of log_rate
-        _callback(is_multiple, _update_tqdm, (value, log_rate))
-
-        # Update progress bar, if remainder
-        _callback(is_remainder, _update_tqdm, (value, remainder))
-
-        # Close progress bar, if last iteration
-        _callback(is_last, _close_tqdm, None)
-
-    def _progress_bar(body_fun: Callable) -> Callable:
-        r"""Decorator that adds a progress bar to `body_fun` used in `jax.lax.scan`."""
-
-        def wrapper_progress_bar(carry: Any, x: Union[tuple, int]) -> Any:
-            # Get iteration number
-            if type(x) is tuple:
-                iter_num, *_ = x
-            else:
-                iter_num = x
-
-            # Compute iteration step
-            result = body_fun(carry, x)
-
-            # Get value
-            *_, value = result
-
-            # Update progress bar
-            _update_progress_bar(value, iter_num)
-
-            return result
-
-        return wrapper_progress_bar
-
-    return _progress_bar
-
-
-__all__ = [
-    "progress_bar",
-]

gpjax-0.9.0.dist-info/LICENSE

@@ -1,201 +0,0 @@
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "{}"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright {yyyy} {name of copyright owner}
-
-   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.