diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 578c6373..2d45beb6 100644 --- a/src/array_api_extra/__init__.py +++ b/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 diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 0edcc7cd..07ddf310 100644 --- a/src/array_api_extra/_delegation.py +++ b/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, /, diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 7f28cbd2..6c73e452 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/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; diff --git a/tests/test_funcs.py b/tests/test_funcs.py index dae7f172..8f52ef1f 100644 --- a/tests/test_funcs.py +++ b/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))