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

gpjax/__init__.py

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

gpjax/citation.py

@@ -8,7 +8,12 @@ from beartype.typing import (
     Dict,
     Union,
 )
-from jaxlib.xla_extension import PjitFunction
+
+try:
+    # safely removable once jax>=0.6.0
+    from jaxlib.xla_extension import PjitFunction
+except ModuleNotFoundError:
+    from jaxlib._jax import PjitFunction
 
 from gpjax.kernels import (
     RFF,
@@ -45,7 +50,7 @@ class AbstractCitation:
 
 
 class NullCitation(AbstractCitation):
-    def __str__(self) -> str:
+    def as_str(self) -> str:
         return (
             "No citation available. If you think this is an error, please open a pull"
             " request."

gpjax/fit.py

@@ -15,13 +15,13 @@
 
 import typing as tp
 
+from flax import nnx
 import jax
+from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random as jr
-import optax as ox
-from flax import nnx
-from jax.flatten_util import ravel_pytree
 from numpyro.distributions.transforms import Transform
+import optax as ox
 from scipy.optimize import minimize
 
 from gpjax.dataset import Dataset

{gpjax-0.11.1.dist-info → gpjax-0.11.2.dist-info}/METADATA

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

{gpjax-0.11.1.dist-info → gpjax-0.11.2.dist-info}/RECORD

@@ -1,8 +1,8 @@
-gpjax/__init__.py,sha256=TjAAfeZTCEl_zsibA8pV76M1jcHkeFhNfWk_SllfgHY,1686
-gpjax/citation.py,sha256=f2Hzj5MLyCE7l0hHAzsEQoTORZH5hgV_eis4uoBiWvE,3811
+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=7L2veA6aRNiozZD8fWa-MVDoYFUKjGJahmvjz8Wp-P0,15046
+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
@@ -43,7 +43,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 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.1.dist-info/METADATA,sha256=02crI6D0dsht6XJ8N1ZqNj5ZktmS5NymVfY45pPmEgM,8558
-gpjax-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.11.1.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.11.1.dist-info/RECORD,,
+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,,