wip

Author: lucascolley

src/array_api_extra/_funcs.py

@@ -292,6 +292,89 @@ def expand_dims(
     return a
 
 
+def isin(
+    element: Array,
+    test_elements: Array,
+    /,
+    *,
+    assume_unique: bool = False,
+    invert: bool = False,
+    xp: ModuleType,
+) -> Array:
+    """Calculates ``element in test_elements``, broadcasting over `element`
+    only.
+
+    Returns a boolean array of the same shape as `element` that is True
+    where an element of `element` is in `test_elements` and False otherwise.
+    """
+
+    original_element_shape = element.shape
+    element = xp.reshape(element, (-1,))
+    test_elements = xp.reshape(test_elements, (-1,))
+    return xp.reshape(
+        _in1d(
+            element,
+            test_elements,
+            assume_unique=assume_unique,
+            invert=invert,
+            xp=xp,
+        ),
+        original_element_shape,
+    )
+
+
+def _in1d(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    assume_unique: bool = False,
+    invert: bool = False,
+    xp: ModuleType,
+) -> Array:
+    """Checks whether each element of an array is also present in a
+    second array.
+
+    Returns a boolean array the same length as `x1` that is True
+    where an element of `x1` is in `x2` and False otherwise.
+
+    This function has been adapted using the original implementation
+    present in numpy:
+    https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
+    """
+
+    # This code is run to make the code significantly faster
+    if x2.shape[0] < 10 * x1.shape[0] ** 0.145:
+        if invert:
+            mask = xp.ones(x1.shape[0], dtype=xp.bool, device=x1.device)
+            for a in x2:
+                mask &= x1 != a
+        else:
+            mask = xp.zeros(x1.shape[0], dtype=xp.bool, device=x1.device)
+            for a in x2:
+                mask |= x1 == a
+        return mask
+
+    if not assume_unique:
+        x1, rev_idx = xp.unique_inverse(x1)
+        x2 = xp.unique_values(x2)
+
+    # ar = xp.concat((x1, x2))
+    # device_ = device(ar)
+    # # We need this to be a stable sort.
+    # order = xp.argsort(ar, stable=True)
+    # reverse_order = xp.argsort(order, stable=True)
+    # sar = xp.take(ar, order, axis=0)
+    # bool_ar = sar[1:] != sar[:-1] if invert else sar[1:] == sar[:-1]
+    # flag = xp.concat((bool_ar, xp.asarray([invert], device=device_)))
+    # ret = xp.take(flag, reverse_order, axis=0)
+
+    # if assume_unique:
+    #     return ret[: x1.shape[0]]
+    # return xp.take(ret, rev_idx, axis=0)
+    return None
+
+
 def kron(a: Array, b: Array, /, *, xp: ModuleType) -> Array:
     """
     Kronecker product of two arrays.
@@ -399,6 +482,22 @@ def kron(a: Array, b: Array, /, *, xp: ModuleType) -> Array:
     return xp.reshape(result, tuple(xp.multiply(a_shape, b_shape)))
 
 
+def setdiff1d(
+    x1: Array, x2: Array, /, *, assume_unique: bool = False, xp: ModuleType
+) -> Array:
+    """Find the set difference of two arrays.
+
+    Return the unique values in `x1` that are not in `x2`.
+    """
+
+    if assume_unique:
+        x1 = xp.reshape(x1, (-1,))
+    else:
+        x1 = xp.unique_values(x1)
+        x2 = xp.unique_values(x2)
+    return x1[_in1d(x1, x2, assume_unique=True, invert=True, xp=xp)]
+
+
 def sinc(x: Array, /, *, xp: ModuleType) -> Array:
     r"""
     Return the normalized sinc function.