gpjax 0.11.1__tar.gz → 0.11.2__tar.gz

{gpjax-0.11.1 → gpjax-0.11.2}/PKG-INFO

@@ -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 → gpjax-0.11.2}/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-0.11.1 → gpjax-0.11.2}/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-0.11.1 → gpjax-0.11.2}/tests/test_fit.py

@@ -13,13 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 
+from beartype.typing import Any
+from flax import nnx
 import jax.numpy as jnp
 import jax.random as jr
+from jaxtyping import (
+    Float,
+    Num,
+)
 import optax as ox
 import pytest
 import scipy
-from beartype.typing import Any
-from flax import nnx
+
 from gpjax.dataset import Dataset
 from gpjax.fit import (
     _check_batch_size,
@@ -54,10 +59,6 @@ from gpjax.parameters import (
 )
 from gpjax.typing import Array
 from gpjax.variational_families import VariationalGaussian
-from jaxtyping import (
-    Float,
-    Num,
-)
 
 
 def test_fit_simple() -> None: