ENH: add nunique delegation

src/array_api_extra/__init__.py

@@ -14,6 +14,7 @@
     kron,
     nan_to_num,
     nanmin,
+    nunique,
     one_hot,
     pad,
     partition,
@@ -30,7 +31,6 @@
     angle,
     apply_where,
     default_dtype,
-    nunique,
 )
 from ._lib._lazy import lazy_apply
 

src/array_api_extra/_delegation.py

@@ -9,14 +9,17 @@
     array_namespace,
     is_cupy_namespace,
     is_dask_namespace,
+    is_jax_array,
     is_jax_namespace,
     is_numpy_namespace,
     is_pydata_sparse_namespace,
     is_torch_namespace,
+    size,
 )
 from ._lib._utils._compat import device as get_device
 from ._lib._utils._helpers import (
     asarrays,
+    capabilities,
     deprecated,
     eager_shape,
     normalize_pad_width,
@@ -36,6 +39,7 @@
     "kron",
     "nan_to_num",
     "nanmin",
+    "nunique",
     "one_hot",
     "pad",
     "partition",
@@ -729,6 +733,48 @@ def nan_to_num(
     return _funcs.nan_to_num(y, fill_value=fill_value, xp=xp)
 
 
+def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
+    """
+    Count the number of unique elements in an array.
+
+    Compatible with JAX and Dask, whose laziness would be otherwise
+    problematic.
+
+    Parameters
+    ----------
+    x : Array
+        Input array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array: 0-dimensional integer array
+        The number of unique elements in `x`. It can be lazy.
+    """
+    if xp is None:
+        xp = array_namespace(x)
+
+    if is_jax_array(x):
+        # size= is JAX-specific
+        # https://github.com/data-apis/array-api/issues/883
+        _, counts = xp.unique_counts(x, size=size(x))
+        return (counts > 0).sum()
+
+    if is_numpy_namespace(xp) or is_cupy_namespace(xp):
+        values = xp.unique(x, equal_nan=False)
+        return xp.asarray(eager_shape(values)[0], device=get_device(x))
+
+    if (
+        is_torch_namespace(xp)
+        and capabilities(xp, device=get_device(x))["data-dependent shapes"]
+    ):
+        values = xp.unique(x)
+        return xp.asarray(eager_shape(values)[0], device=get_device(x))
+
+    return _funcs.nunique(x, xp=xp)
+
+
 def one_hot(
     x: Array,
     /,

src/array_api_extra/_lib/_funcs.py

@@ -11,7 +11,6 @@
 from ._utils._compat import (
     array_namespace,
     is_dask_namespace,
-    is_jax_array,
 )
 from ._utils._helpers import (
     asarrays,
@@ -550,34 +549,8 @@ def perform_replacements(  # numpydoc ignore=PR01,RT01
     return x
 
 
-def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
-    """
-    Count the number of unique elements in an array.
-
-    Compatible with JAX and Dask, whose laziness would be otherwise
-    problematic.
-
-    Parameters
-    ----------
-    x : Array
-        Input array.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `x`. Default: infer.
-
-    Returns
-    -------
-    array: 0-dimensional integer array
-        The number of unique elements in `x`. It can be lazy.
-    """
-    if xp is None:
-        xp = array_namespace(x)
-
-    if is_jax_array(x):
-        # size= is JAX-specific
-        # https://github.com/data-apis/array-api/issues/883
-        _, counts = xp.unique_counts(x, size=_compat.size(x))
-        return (counts > 0).sum()
-
+def nunique(x: Array, /, *, xp: ModuleType) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in `array_api_extra._delegation.py`."""
     # There are 3 general use cases:
     # 1. backend has unique_counts and it returns an array with known shape
     # 2. backend has unique_counts and it returns a None-sized array;

tests/test_funcs.py

@@ -1441,6 +1441,27 @@ def test_fill_value_failure(self, xp: ModuleType) -> None:
 
 
 class TestNUnique:
+    @pytest.mark.skip_xp_backend(
+        Backend.ARRAY_API_STRICT, reason="array-agnostic fallback"
+    )
+    @pytest.mark.skip_xp_backend(
+        Backend.ARRAY_API_STRICTEST, reason="array-agnostic fallback"
+    )
+    @pytest.mark.skip_xp_backend(Backend.DASK, reason="array-agnostic fallback")
+    @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="array-agnostic fallback")
+    def test_delegates(
+        self,
+        xp: ModuleType,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        def fallback(*_args: object, **_kwargs: object) -> Array:
+            msg = "array-agnostic fallback should not be used"
+            raise AssertionError(msg)
+
+        monkeypatch.setattr(functions, "nunique", fallback)
+        a = xp.asarray([1, 1, 2])
+        assert_equal(nunique(a), xp.asarray(2))
+
     def test_simple(self, xp: ModuleType):
         a = xp.asarray([[1, 1], [0, 2], [2, 2]])
         assert_equal(nunique(a), xp.asarray(3))