Repository: cupy/cupy

Issue:
Support `axes` option in `cupy.unique`
### Description

Currently `axes` option is `cupy.unique` is unimplemented.

### Additional Information

_No response_


Discussion:

Resolving patch:
diff --git a/cupy/_manipulation/add_remove.py b/cupy/_manipulation/add_remove.py
--- a/cupy/_manipulation/add_remove.py
+++ b/cupy/_manipulation/add_remove.py
@@ -1,6 +1,7 @@
 import numpy
 
 import cupy
+import math
 from cupy import _core
 
 
@@ -194,7 +195,11 @@ def unique(ar, return_index=False, return_inverse=False,
             to reconstruct `ar`.
         return_counts(bool, optional): If True, also return the number of times
             each unique item appears in `ar`.
-        axis(int or None, optional): Not supported yet.
+        axis(int or None, optional): The axis to operate on. If None, ar will
+            be flattened. If an integer, the subarrays indexed by the given
+            axis will be flattened and treated as the elements of a 1-D array
+            with the dimension of the given axis, see the notes for more
+            details. The default is None.
         equal_nan(bool, optional): If True, collapse multiple NaN values in the
             return array into one.
 
@@ -212,15 +217,115 @@ def unique(ar, return_index=False, return_inverse=False,
             * The number of times each of the unique values comes up in the
               original array. Only provided if `return_counts` is True.
 
+    Notes:
+       When an axis is specified the subarrays indexed by the axis are sorted.
+       This is done by making the specified axis the first dimension of the
+       array (move the axis to the first dimension to keep the order of the
+       other axes) and then flattening the subarrays in C order.
+
     .. warning::
 
         This function may synchronize the device.
 
     .. seealso:: :func:`numpy.unique`
     """
-    if axis is not None:
-        raise NotImplementedError('axis option is not supported yet.')
+    if axis is None:
+        ret = _unique_1d(ar, return_index=return_index,
+                         return_inverse=return_inverse,
+                         return_counts=return_counts,
+                         equal_nan=equal_nan)
+        return ret
+
+    ar = cupy.moveaxis(ar, axis, 0)
+
+    # The array is reshaped into a contiguous 2D array
+    orig_shape = ar.shape
+    idx = cupy.arange(0, orig_shape[0], dtype=cupy.intp)
+    ar = ar.reshape(orig_shape[0], math.prod(orig_shape[1:]))
+    ar = cupy.ascontiguousarray(ar)
+    is_unsigned = cupy.issubdtype(ar.dtype, cupy.unsignedinteger)
+    is_complex = cupy.iscomplexobj(ar)
+
+    ar_cmp = ar
+    if is_unsigned:
+        ar_cmp = ar.astype(cupy.intp)
+
+    def compare_axis_elems(idx1, idx2):
+        left, right = ar_cmp[idx1], ar_cmp[idx2]
+        comp = cupy.trim_zeros(left - right, 'f')
+        if comp.shape[0] > 0:
+            diff = comp[0]
+            if is_complex and cupy.isnan(diff):
+                return True
+            return diff < 0
+        return False
+
+    # The array is sorted lexicographically using the first item of each
+    # element on the axis
+    sorted_indices = cupy.empty(orig_shape[0], dtype=cupy.intp)
+    queue = [(idx.tolist(), 0)]
+    while queue != []:
+        current, off = queue.pop(0)
+        if current == []:
+            continue
+
+        mid_elem = current[0]
+        left = []
+        right = []
+        for i in range(1, len(current)):
+            if compare_axis_elems(current[i], mid_elem):
+                left.append(current[i])
+            else:
+                right.append(current[i])
+
+        elem_pos = off + len(left)
+        queue.append((left, off))
+        queue.append((right, elem_pos + 1))
+
+        sorted_indices[elem_pos] = mid_elem
+
+    ar = ar[sorted_indices]
+
+    if ar.size > 0:
+        mask = cupy.empty(ar.shape, dtype=cupy.bool_)
+        mask[:1] = True
+        mask[1:] = ar[1:] != ar[:-1]
+
+        mask = cupy.any(mask, axis=1)
+    else:
+        # If empty, then the mask should grab the first empty array as the
+        # unique one
+        mask = cupy.ones((ar.shape[0]), dtype=cupy.bool_)
+        mask[1:] = False
+
+    # Index the input array with the unique elements and reshape it into the
+    # original size and dimension order
+    ar = ar[mask]
+    ar = ar.reshape(mask.sum().item(), *orig_shape[1:])
+    ar = cupy.moveaxis(ar, 0, axis)
+
+    ret = ar,
+    if return_index:
+        ret += sorted_indices[mask],
+    if return_inverse:
+        imask = cupy.cumsum(mask) - 1
+        inv_idx = cupy.empty(mask.shape, dtype=cupy.intp)
+        inv_idx[sorted_indices] = imask
+        ret += inv_idx,
+    if return_counts:
+        nonzero = cupy.nonzero(mask)[0]  # may synchronize
+        idx = cupy.empty((nonzero.size + 1,), nonzero.dtype)
+        idx[:-1] = nonzero
+        idx[-1] = mask.size
+        ret += idx[1:] - idx[:-1],
+
+    if len(ret) == 1:
+        ret = ret[0]
+    return ret
+
 
+def _unique_1d(ar, return_index=False, return_inverse=False,
+               return_counts=False, equal_nan=True):
     ar = cupy.asarray(ar).flatten()
 
     if return_index or return_inverse:


Test patch:
diff --git a/tests/cupy_tests/manipulation_tests/test_add_remove.py b/tests/cupy_tests/manipulation_tests/test_add_remove.py
--- a/tests/cupy_tests/manipulation_tests/test_add_remove.py
+++ b/tests/cupy_tests/manipulation_tests/test_add_remove.py
@@ -149,53 +149,99 @@ class TestUnique:
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique(self, xp, dtype):
+    def test_unique_no_axis(self, xp, dtype):
         a = testing.shaped_random((100, 100), xp, dtype)
         return xp.unique(a)
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique_index(self, xp, dtype):
+    def test_unique(self, xp, dtype):
+        a = testing.shaped_random((100, 100), xp, dtype)
+        return xp.unique(a, axis=1)
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_index_no_axis(self, xp, dtype):
         a = testing.shaped_random((100, 100), xp, dtype)
         return xp.unique(a, return_index=True)[1]
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique_inverse(self, xp, dtype):
+    def test_unique_index(self, xp, dtype):
+        a = testing.shaped_random((100, 100), xp, dtype)
+        return xp.unique(a, return_index=True, axis=0)[1]
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_inverse_no_axis(self, xp, dtype):
         a = testing.shaped_random((100, 100), xp, dtype)
         return xp.unique(a, return_inverse=True)[1]
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique_counts(self, xp, dtype):
+    def test_unique_inverse(self, xp, dtype):
+        a = testing.shaped_random((100, 100), xp, dtype)
+        return xp.unique(a, return_inverse=True, axis=1)[1]
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_counts_no_axis(self, xp, dtype):
         a = testing.shaped_random((100, 100), xp, dtype)
         return xp.unique(a, return_counts=True)[1]
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique_return_all(self, xp, dtype):
+    def test_unique_counts(self, xp, dtype):
+        a = testing.shaped_random((100, 100), xp, dtype)
+        return xp.unique(a, return_counts=True, axis=0)[1]
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_return_all_no_axis(self, xp, dtype):
         a = testing.shaped_random((100, 100), xp, dtype)
         return xp.unique(
             a, return_index=True, return_inverse=True, return_counts=True)
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique_empty(self, xp, dtype):
+    def test_unique_return_all(self, xp, dtype):
+        a = testing.shaped_random((100, 100), xp, dtype)
+        return xp.unique(
+            a, return_index=True, return_inverse=True, return_counts=True,
+            axis=1)
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_empty_no_axis(self, xp, dtype):
         a = xp.empty((0,), dtype)
         return xp.unique(a)
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
-    def test_unique_empty_return_all(self, xp, dtype):
+    def test_unique_empty(self, xp, dtype):
+        a = xp.empty((0,), dtype)
+        return xp.unique(a, axis=0)
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_empty_return_all_no_axis(self, xp, dtype):
         a = xp.empty((3, 0, 2), dtype)
         return xp.unique(
             a, return_index=True, return_inverse=True, return_counts=True)
 
+    @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
+    @testing.numpy_cupy_array_equal()
+    def test_unique_empty_return_all(self, xp, dtype):
+        a = xp.empty((3, 0, 2), dtype)
+        return xp.unique(
+            a, return_index=True, return_inverse=True, return_counts=True,
+            axis=2)
+
     @pytest.mark.parametrize('equal_nan', [True, False])
     @pytest.mark.parametrize('dtype', 'efdFD')
     @testing.numpy_cupy_array_equal()
     @testing.with_requires('numpy>=1.23.1')
-    def test_unique_equal_nan(self, xp, dtype, equal_nan):
+    def test_unique_equal_nan_no_axis(self, xp, dtype, equal_nan):
         if xp.dtype(dtype).kind == 'c':
             # Nan and Nan+Nan*1j are collapsed when equal_nan=True
             a = xp.array([
@@ -206,6 +252,26 @@ def test_unique_equal_nan(self, xp, dtype, equal_nan):
             a = xp.array([2, xp.nan, 2, xp.nan, 1], dtype=dtype)
         return xp.unique(a, equal_nan=equal_nan)
 
+    @pytest.mark.parametrize('equal_nan', [True, False])
+    @pytest.mark.parametrize('dtype', 'fdFD')
+    @testing.numpy_cupy_array_equal()
+    @testing.with_requires('numpy>=1.23.1')
+    def test_unique_equal_nan(self, xp, dtype, equal_nan):
+        if xp.dtype(dtype).kind == 'c':
+            # Nan and Nan+Nan*1j are collapsed when equal_nan=True
+            a = xp.array([
+                [complex(xp.nan, 3), 2, complex(7, xp.nan)],
+                [xp.nan, complex(xp.nan, xp.nan), 2],
+                [xp.nan, 1, complex(xp.nan, -1)]
+            ], dtype=dtype)
+        else:
+            a = xp.array([
+                [2, xp.nan, 2],
+                [xp.nan, 1, xp.nan],
+                [xp.nan, 1, xp.nan]
+            ], dtype=dtype)
+        return xp.unique(a, axis=0, equal_nan=equal_nan)
+
 
 @testing.parameterize(*testing.product({
     'trim': ['fb', 'f', 'b']