Ensure that Jax float0 array is recognized

See jax-ml/jax#20620.

Author: NeilGirdhar

array_api_compat/common/_helpers.py

@@ -18,6 +18,20 @@
 import inspect
 import warnings
 
+def is_jax_zero_gradient_array(x):
+    """Return True if `x` is a zero-gradient array.
+
+    These arrays are a design quirk of Jax that may one day be removed.
+    See https://github.com/google/jax/issues/20620.
+    """
+    if 'numpy' not in sys.modules or 'jax' not in sys.modules:
+        return False
+
+    import numpy as np
+    import jax
+
+    return isinstance(x, np.ndarray) and x.dtype == jax.float0
+
 def is_numpy_array(x):
     """
     Return True if `x` is a NumPy array.
@@ -44,7 +58,8 @@ def is_numpy_array(x):
     import numpy as np
 
     # TODO: Should we reject ndarray subclasses?
-    return isinstance(x, (np.ndarray, np.generic))
+    return (isinstance(x, (np.ndarray, np.generic))
+            and not is_jax_zero_gradient_array(x))
 
 def is_cupy_array(x):
     """
@@ -149,7 +164,7 @@ def is_jax_array(x):
 
     import jax
 
-    return isinstance(x, jax.Array)
+    return isinstance(x, jax.Array) or is_jax_zero_gradient_array(x)
 
 def is_array_api_obj(x):
     """