WIP ENH: in1d jax.jit support

Author: crusaderky

pixi.lock

@@ -230,6 +230,7 @@ reportUnknownVariableType = false
 # Redundant with mypy checks
 reportMissingImports = false
 reportMissingTypeStubs = false
+reportPossiblyUnboundVariable = false
 # false positives for input validation
 reportUnreachable = false
 # ruff handles this

pyproject.toml

@@ -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/_funcs.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"]
@@ -45,13 +46,15 @@ def in1d(
                 mask |= x1 == a
         return mask
 
-    rev_idx = xp.empty(0)  # placeholder
     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)
@@ -62,12 +65,12 @@ 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 assume_unique:
-        return ret[: x1.shape[0]]
-    return xp.take(ret, rev_idx, axis=0)
+    if not assume_unique:
+        return xp.take(ret, rev_idx, axis=0)  # type: ignore[possibly-undefined]
+    return ret[: x1.shape[0]]
 
 
 def mean(

src/array_api_extra/_lib/_utils/_helpers.py

@@ -34,7 +34,7 @@
 lazy_xp_function(kron, static_argnames="xp")
 lazy_xp_function(nunique, static_argnames="xp")
 lazy_xp_function(pad, static_argnames=("pad_width", "mode", "constant_values", "xp"))
-# FIXME calls in1d which calls xp.unique_values without size
+lazy_xp_function(setdiff1d, static_argnames=("assume_unique", "xp"))
 lazy_xp_function(setdiff1d, jax_jit=False, static_argnames=("assume_unique", "xp"))
 lazy_xp_function(sinc, static_argnames="xp")
 

tests/test_funcs.py

@@ -11,8 +11,7 @@
 
 # mypy: disable-error-code=no-untyped-usage
 
-# FIXME calls xp.unique_values without size
-lazy_xp_function(in1d, jax_jit=False, static_argnames=("assume_unique", "invert", "xp"))
+lazy_xp_function(in1d, static_argnames=("assume_unique", "invert", "xp"))
 
 
 class TestIn1D: