fix: jax 4.3.1 removal of FInfo (#64)

* fix: jax 4.3.1 removal of FInfo * fix: lax.convert_element_type_p any kwargs Signed-off-by: nstarman <[email protected]>
GalacticDynamics · Jul 30, 2024 · 6de6552 · 6de6552
1 parent 0a88ecf
commit 6de6552
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 12 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -28,8 +28,8 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "jax != 0.4.30",
-    "jaxlib != 0.4.30",
+    "jax > 0.4, != 0.4.30",
+    "jaxlib > 0.4, != 0.4.30",
     "plum-dispatch",
     "quax>=0.0.3",
 ]

diff --git a/src/quaxed/array_api/_data_type_functions.py b/src/quaxed/array_api/_data_type_functions.py
@@ -2,12 +2,17 @@
 
 
 from jax.experimental import array_api
-from jax.experimental.array_api._data_type_functions import FInfo
 from jaxtyping import ArrayLike
 
 from quaxed._types import DType
 from quaxed._utils import JAX_VERSION, quaxify
 
+if JAX_VERSION < (0, 4, 31):
+    from jax.experimental.array_api._data_type_functions import FInfo
+else:
+    from numpy import finfo as FInfo  # noqa: N812
+
+
 if JAX_VERSION < (0, 4, 27):
     from jax.experimental.array_api._data_type_functions import IInfo
 else:

diff --git a/tests/array_api/test_myarray.py b/tests/array_api/test_myarray.py
@@ -225,7 +225,7 @@ def test_finfo():
     expected = jax_xp.finfo(jnp.float32)
 
     assert isinstance(got, FInfo)
-    for attr in FInfo.__slots__:
+    for attr in ("bits", "eps", "max", "min", "smallest_normal", "dtype"):
         assert getattr(got, attr) == getattr(expected, attr)
 
 

diff --git a/tests/myarray.py b/tests/myarray.py
@@ -305,14 +305,11 @@ def _conv_general_dilated_p() -> MyArray:
 
 
 @register(lax.convert_element_type_p)
-def _convert_element_type_p(
-    operand: MyArray,
-    *,
-    new_dtype: Any,
-    weak_type: Any,
-) -> MyArray:
-    del weak_type
-    return replace(operand, array=lax.convert_element_type(operand.array, new_dtype))
+def _convert_element_type_p(operand: MyArray, **kwargs: Any) -> MyArray:
+    return replace(
+        operand,
+        array=lax.convert_element_type_p.bind(operand.array, **kwargs),
+    )
 
 
 # ==============================================================================