Support for multi-dimensional arrays

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -331,6 +331,8 @@ def pad(
 def partition(
     a: Array,
     kth: int,
+    /,
+    axis: int | None = -1,
     *,
     xp: ModuleType | None = None,
 ) -> Array:
@@ -343,6 +345,9 @@ def partition(
         Input array.
     kth : int
         Element index to partition by.
+    axis : int, optional
+        Axis along which to partition. The default is -1 (the last axis).
+        If None, the flattened array is used.
     xp : array_namespace, optional
         The standard-compatible namespace for `x`. Default: infer.
 
@@ -354,36 +359,61 @@ def partition(
     # Validate inputs.
     if xp is None:
         xp = array_namespace(a)
-    if a.ndim != 1:
-        msg = "only 1-dimensional arrays are currently supported"
-        raise NotImplementedError(msg)
+    if a.ndim < 1:
+        msg = "`a` must be at least 1-dimensional"
+        raise TypeError(msg)
+    if axis is None:
+        return partition(xp.reshape(a, -1), kth, axis=0, xp=xp)
+    size = a.shape[axis]
+    if size is None:
+        msg = "Array dimensions must be known"
+        raise ValueError(msg)
+    if not (0 <= kth < size):
+        msg = f"kth(={kth}) out of bounds [0 {size})"
+        raise ValueError(msg)
 
     # Delegate where possible.
-    if is_numpy_namespace(xp) or is_cupy_namespace(xp):
-        return xp.partition(a, kth)
-    if is_jax_namespace(xp):
-        from jax import numpy
-
-        return numpy.partition(a, kth)
+    if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp):
+        return xp.partition(a, kth, axis=axis)
 
     # Use top-k when possible:
     if is_torch_namespace(xp):
-        from torch import topk
+        if not (axis == -1 or axis == a.ndim - 1):
+            a = xp.transpose(a, axis, -1)
 
-        a_left, indices_left = topk(a, kth, largest=False, sorted=False)
+        # Get smallest `kth` elements along axis
+        kth += 1  # HACK: we use a non-specified behavior of torch.topk:
+        # in `a_left`, the element in the last position is the max
+        a_left, indices = xp.topk(a, kth, dim=-1, largest=False, sorted=False)
+
+        # Build a mask to remove the selected elements
         mask_right = xp.ones(a.shape, dtype=bool)
-        mask_right[indices_left] = False
-        return xp.concat((a_left, a[mask_right]))
+        mask_right.scatter_(dim=-1, index=indices, value=False)
+
+        # Remaining elements along axis
+        a_right = a[mask_right]  # 1-d array
+
+        # Reshape. This is valid only because we work on the last axis
+        a_right = xp.reshape(a_right, shape=(*a.shape[:-1], -1))
+
+        # Concatenate the two parts along axis
+        partitioned_array = xp.cat((a_left, a_right), dim=-1)
+        if not (axis == -1 or axis == a.ndim - 1):
+            partitioned_array = xp.transpose(partitioned_array, axis, -1)
+        return partitioned_array
+
     # Note: dask topk/argtopk sort the return values, so it's
     # not much more efficient than sorting everything when
     # kth is not small compared to x.size
 
-    return _funcs.partition(a, kth, xp=xp)
+    return _funcs.partition(a, kth, axis=axis, xp=xp)
 
 
 def argpartition(
     a: Array,
     kth: int,
+    /,
+    axis: int | None = -1,
     *,
     xp: ModuleType | None = None,
 ) -> Array:
@@ -392,10 +422,13 @@ def argpartition(
 
     Parameters
     ----------
-    a : 1-dimensional array
+    a : Array
         Input array.
     kth : int
         Element index to partition by.
+    axis : int, optional
+        Axis along which to partition. The default is -1 (the last axis).
+        If None, the flattened array is used.
     xp : array_namespace, optional
         The standard-compatible namespace for `x`. Default: infer.
 
@@ -407,29 +440,46 @@ def argpartition(
     # Validate inputs.
     if xp is None:
         xp = array_namespace(a)
-    if a.ndim != 1:
-        msg = "only 1-dimensional arrays are currently supported"
-        raise NotImplementedError(msg)
+    if a.ndim < 1:
+        msg = "`a` must be at least 1-dimensional"
+        raise TypeError(msg)
+    if axis is None:
+        return partition(xp.reshape(a, -1), kth, axis=0, xp=xp)
+    size = a.shape[axis]
+    if size is None:
+        msg = "Array dimensions must be known"
+        raise ValueError(msg)
+    if not (0 <= kth < size):
+        msg = f"kth(={kth}) out of bounds [0 {size})"
+        raise ValueError(msg)
 
     # Delegate where possible.
-    if is_numpy_namespace(xp) or is_cupy_namespace(xp):
-        return xp.argpartition(a, kth)
-    if is_jax_namespace(xp):
-        from jax import numpy
-
-        return numpy.argpartition(a, kth)
+    if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp):
+        return xp.argpartition(a, kth, axis=axis)
 
     # Use top-k when possible:
     if is_torch_namespace(xp):
-        from torch import topk
+        # see `partition` above for commented details of those steps:
+        if not (axis == -1 or axis == a.ndim - 1):
+            a = xp.transpose(a, axis, -1)
+
+        kth += 1  # HACK
+        _, indices_left = xp.topk(a, kth, dim=-1, largest=False, sorted=False)
+
+        mask_right = xp.ones(a.shape, dtype=bool)
+        mask_right.scatter_(dim=-1, index=indices_left, value=False)
+
+        indices_right = xp.nonzero(mask_right)[-1]
+        indices_right = xp.reshape(indices_right, shape=(*a.shape[:-1], -1))
+
+        # Concatenate the two parts along axis
+        index_array = xp.cat((indices_left, indices_right), dim=-1)
+        if not (axis == -1 or axis == a.ndim - 1):
+            index_array = xp.transpose(index_array, axis, -1)
+        return index_array
 
-        _, indices = topk(a, kth, largest=False, sorted=False)
-        mask = xp.ones(a.shape, dtype=bool)
-        mask[indices] = False
-        indices_above = xp.arange(a.shape[0])[mask]
-        return xp.concat((indices, indices_above))
     # Note: dask topk/argtopk sort the return values, so it's
     # not much more efficient than sorting everything when
     # kth is not small compared to x.size
 
-    return _funcs.argpartition(a, kth, xp=xp)
+    return _funcs.argpartition(a, kth, axis=axis, xp=xp)

src/array_api_extra/_lib/_funcs.py

@@ -1034,18 +1034,22 @@ def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
 def partition(  # numpydoc ignore=PR01,RT01
     x: Array,
     kth: int,  # noqa: ARG001
+    /,
+    axis: int = -1,
     *,
     xp: ModuleType,
 ) -> Array:
     """See docstring in `array_api_extra._delegation.py`."""
-    return xp.sort(x, stable=False)
+    return xp.sort(x, axis=axis, stable=False)
 
 
 def argpartition(  # numpydoc ignore=PR01,RT01
     x: Array,
     kth: int,  # noqa: ARG001
+    /,
+    axis: int = -1,
     *,
     xp: ModuleType,
 ) -> Array:
     """See docstring in `array_api_extra._delegation.py`."""
-    return xp.argsort(x, stable=False)
+    return xp.argsort(x, axis=axis, stable=False)

tests/test_funcs.py

@@ -9,6 +9,7 @@
 import pytest
 from hypothesis import given
 from hypothesis import strategies as st
+from typing_extensions import override
 
 from array_api_extra import (
     apply_where,
@@ -32,7 +33,12 @@
 )
 from array_api_extra._lib._backends import NUMPY_VERSION, Backend
 from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
-from array_api_extra._lib._utils._compat import device as get_device
+from array_api_extra._lib._utils._compat import (
+    device as get_device,
+)
+from array_api_extra._lib._utils._compat import (
+    is_pydata_sparse_namespace,
+)
 from array_api_extra._lib._utils._helpers import eager_shape, ndindex
 from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import lazy_xp_function
@@ -1303,15 +1309,67 @@ def test_xp(self, xp: ModuleType):
 
 
 class TestPartition:
-    def test_basic(self, xp: ModuleType):
-        # Using 0-dimensional array
-        rng = np.random.default_rng(2847)
-
-        for _ in range(100):
-            n = rng.integers(1, 1000)
-            x = xp.asarray(rng.random(size=n))
-            k = int(rng.integers(1, n - 1))
-            y = partition(x, k)
-            assert xp.max(y[:k]) <= xp.min(y[k:])
-            y = x[argpartition(x, k)]
-            assert xp.max(y[:k]) <= xp.min(y[k:])
+    @classmethod
+    def _assert_valid_partition(cls, x: Array, k: int, xp: ModuleType, axis: int = -1):
+        if x.ndim != 1 and axis == 0:
+            assert isinstance(x.shape[1], int)
+            for i in range(x.shape[1]):
+                cls._assert_valid_partition(x[:, i, ...], k, xp, axis=0)
+        elif x.ndim != 1:
+            axis = axis - 1 if axis != -1 else -1
+            assert isinstance(x.shape[0], int)
+            for i in range(x.shape[0]):
+                cls._assert_valid_partition(x[i, ...], k, xp, axis=axis)
+        else:
+            if k > 0:
+                assert xp.max(x[:k]) <= x[k]
+            assert x[k] <= xp.min(x[k:])
+
+    @classmethod
+    def _partition(
+        cls,
+        x: Array,
+        k: int,
+        xp: ModuleType,  # noqa: ARG003
+        axis: int | None = -1,
+    ):
+        return partition(x, k, axis=axis)
+
+    def test_1d(self, xp: ModuleType):
+        rng = np.random.default_rng()
+        for n in [2, 3, 4, 5, 7, 10, 20, 50, 100, 1_000]:
+            k = int(rng.integers(n))
+            x = xp.asarray(rng.integers(n, size=n))
+            self._assert_valid_partition(self._partition(x, k, xp), k, xp)
+            x = xp.asarray(rng.random(n))
+            self._assert_valid_partition(self._partition(x, k, xp), k, xp)
+
+    @pytest.mark.parametrize("ndim", [2, 3, 4, 5])
+    def test_nd(self, xp: ModuleType, ndim: int):
+        rng = np.random.default_rng()
+
+        for n in [2, 3, 5, 10, 20, 100]:
+            base_shape = [int(v) for v in rng.integers(1, 4, size=ndim)]
+            k = int(rng.integers(n))
+
+            for i in range(ndim):
+                shape = base_shape[:]
+                shape[i] = n
+                x = xp.asarray(rng.integers(n, size=tuple(shape)))
+                y = self._partition(x, k, xp, axis=i)
+                self._assert_valid_partition(y, k, xp, axis=i)
+
+
+@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort")
+class TestArgpartition(TestPartition):
+    @classmethod
+    @override
+    def _partition(cls, x: Array, k: int, xp: ModuleType, axis: int | None = -1):
+        if is_pydata_sparse_namespace(xp):
+            pytest.xfail(reason="Sparse backend has no argsort")
+        indices = argpartition(x, k, axis=axis)
+        if x.ndim == 1:
+            return x[indices]
+        if not hasattr(xp, "take_along_axis"):
+            pytest.skip("TODO: find an alternative to take_along_axis")
+        return xp.take_along_axis(x, indices, axis=axis)