WIP

Author: crusaderky

src/array_api_extra/_lib/_funcs.py

@@ -11,7 +11,7 @@
 
 from ._at import at
 from ._utils import _compat, _helpers
-from ._utils._compat import array_namespace, is_jax_array
+from ._utils._compat import array_namespace, is_jax_array, is_jax_namespace
 from ._utils._typing import Array
 
 __all__ = [
@@ -533,6 +533,29 @@ def pad(
     return at(padded, tuple(slices)).set(x)
 
 
+def _ensure_unique_values(x: Array, assume_unique: bool, xp: ModuleType) -> Array:
+    """
+    Wrapper around xp.unique_values
+
+    If x is a JAX array and we're running inside jax.jit, the output
+    shape needs to be known. Return an array the size of x, padded
+    with the first element of x.
+    """
+    x = xp.reshape(x, (-1,))
+    if assume_unique or x.shape == (0, ):
+        return x
+
+    if is_jax_array(x):
+        import jax
+
+        try:
+            return xp.unique_values(x)  # eager
+        except jax.errors.ConcretizationError:  # inside jax.jit
+            return xp.unique_values(x, size=x.size, fill_value=x[0])
+
+    return xp.unique_values(x)
+
+
 def setdiff1d(
     x1: Array,
     x2: Array,
@@ -578,12 +601,17 @@ def setdiff1d(
     if xp is None:
         xp = array_namespace(x1, x2)
 
-    if assume_unique:
-        x1 = xp.reshape(x1, (-1,))
-    else:
-        x1 = xp.unique_values(x1)
-        x2 = xp.unique_values(x2)
-    return x1[_helpers.in1d(x1, x2, assume_unique=True, invert=True, xp=xp)]
+    x1 = _ensure_unique_values(x1, assume_unique=assume_unique, xp=xp)
+    x2 = _ensure_unique_values(x2, assume_unique=assume_unique, xp=xp)
+
+    if x1.shape == (0,) or x2.shape == (0, ):
+        return x1
+
+    mask = _helpers.in1d(x1, x2, assume_unique=True, invert=True, xp=xp)
+    if is_jax_namespace(xp):
+        fill = xp.zeros((), dtype=x1.dtype, device=_compat.device(x1))
+        return xp.where(mask, fill, x1)
+    return x1[mask]
 
 
 def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:

src/array_api_extra/_lib/_utils/_helpers.py

@@ -6,6 +6,7 @@
 from types import ModuleType
 
 from . import _compat
+from ._compat import is_jax_namespace
 from ._typing import Array
 
 __all__ = ["in1d", "mean"]
@@ -46,11 +47,14 @@ def in1d(
         return mask
 
     if not assume_unique:
-        x1, rev_idx = xp.unique_inverse(x1)
-        x2 = xp.unique_values(x2)
+        if is_jax_namespace(xp):
+            x1, rev_idx = xp.unique_inverse(x1, size=x1.size, fill_value=x1[0])
+            x2 = xp.unique_values(x2, size=x2.size, fill_value=x2[0])
+        else:
+            x1, rev_idx = xp.unique_inverse(x1)
+            x2 = xp.unique_values(x2)
 
     ar = xp.concat((x1, x2))
-    device_ = _compat.device(ar)
     # We need this to be a stable sort.
     order = xp.argsort(ar, stable=True)
     reverse_order = xp.argsort(order, stable=True)
@@ -61,7 +65,7 @@ def in1d(
         bool_ar = sar[1:] != sar[:-1] if invert else sar[1:] == sar[:-1]
     else:
         bool_ar = xp.asarray([False]) if invert else xp.asarray([True])
-    flag = xp.concat((bool_ar, xp.asarray([invert], device=device_)))
+    flag = xp.concat((bool_ar, xp.asarray([invert], device=_compat.device(ar))))
     ret = xp.take(flag, reverse_order, axis=0)
 
     if not assume_unique: