build: fix jax bump (#51)

* build: fix jax bump Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · May 7, 2024 · 7bfacbe · 7bfacbe
1 parent e796b50
commit 7bfacbe
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/src/quaxed/_utils.py b/src/quaxed/_utils.py
@@ -1,3 +1,4 @@
+from importlib.metadata import version
 from typing import TypeVar
 
 import quax
@@ -8,3 +9,6 @@
 def quaxify(func: T, *, filter_spec: bool | tuple[bool, ...] = True) -> T:
     """Quaxify, but makes mypy happy."""
     return quax.quaxify(func, filter_spec=filter_spec)
+
+
+JAX_VERSION: tuple[int, ...] = tuple(map(int, version("jax").split(".")))
diff --git a/src/quaxed/array_api/_data_type_functions.py b/src/quaxed/array_api/_data_type_functions.py
@@ -2,11 +2,16 @@
 
 
 from jax.experimental import array_api
-from jax.experimental.array_api._data_type_functions import FInfo, IInfo
+from jax.experimental.array_api._data_type_functions import FInfo
 from jaxtyping import ArrayLike
 
 from quaxed._types import DType
-from quaxed._utils import quaxify
+from quaxed._utils import JAX_VERSION, quaxify
+
+if JAX_VERSION < (0, 4, 27):
+    from jax.experimental.array_api._data_type_functions import IInfo
+else:
+    from jax.experimental.array_api import iinfo as IInfo  # noqa: N812
 
 
 @quaxify

diff --git a/tests/array_api/test_myarray.py b/tests/array_api/test_myarray.py
@@ -9,9 +9,9 @@
     _UniqueCountsResult,
     _UniqueInverseResult,
 )
-from jax.experimental.array_api._data_type_functions import FInfo, IInfo
 
 import quaxed.array_api as xp
+from quaxed.array_api._data_type_functions import FInfo, IInfo
 
 from ..myarray import MyArray
 
@@ -235,7 +235,7 @@ def test_iinfo():
     expected = jax_xp.iinfo(jnp.int32)
 
     assert isinstance(got, IInfo)
-    for attr in IInfo.__slots__:
+    for attr in ("kind", "bits", "min", "max", "dtype"):
         assert getattr(got, attr) == getattr(expected, attr)