add xfail according to jax issue

Author: Enderdead

src/array_api_extra/_delegation.py

@@ -462,8 +462,10 @@ def setdiff1d(
     if xp is None:
         xp = array_namespace(x1, x2)
 
-    if is_numpy_namespace(xp) or is_jax_namespace(xp) or is_cupy_namespace(xp):
-        return xp.setdiff1d(x1, x2, assume_unique=assume_unique)
+    if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp):
+        # https://github.com/microsoft/pyright/issues/10103
+        x1_, x2_ = asarrays(x1, x2, xp=xp)
+        return xp.setdiff1d(x1_, x2_, assume_unique=assume_unique)
 
     return _funcs.setdiff1d(x1, x2, assume_unique=assume_unique, xp=xp)
 

tests/test_funcs.py

@@ -30,7 +30,7 @@
 )
 from array_api_extra._lib._backends import NUMPY_VERSION, Backend
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
-from array_api_extra._lib._utils._compat import device as get_device
+from array_api_extra._lib._utils._compat import device as get_device, is_jax_namespace
 from array_api_extra._lib._utils._helpers import eager_shape, ndindex
 from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import lazy_xp_function
@@ -1228,6 +1228,10 @@ def test_shapes(
     ):
         x1 = xp.zeros(shape1)
         x2 = xp.zeros(shape2)
+
+        if is_jax_namespace(xp) and assume_unique and shape1!=(1,):
+            pytest.xfail(reason="jax#32335 fixed with jax>=0.8.0")
+
         actual = setdiff1d(x1, x2, assume_unique=assume_unique)
         xp_assert_equal(actual, xp.empty((0,)))
 
@@ -1240,6 +1244,9 @@ def test_python_scalar(self, xp: ModuleType, assume_unique: bool):
         actual = setdiff1d(x1, x2, assume_unique=assume_unique)
         xp_assert_equal(actual, xp.asarray([1, 2], dtype=xp.int16))
 
+        if is_jax_namespace(xp) and assume_unique:
+            pytest.xfail(reason="jax#32335 fixed with jax>=0.8.0")
+
         actual = setdiff1d(x2, x1, assume_unique=assume_unique)
         xp_assert_equal(actual, xp.asarray([], dtype=xp.int16))