dpnp.unique returns reshaped unique_inverse when axis is None

dpnp/dpnp_iface_manipulation.py

@@ -2229,6 +2229,12 @@ def unique(
     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.
+    For complex arrays all NaN values are considered equivalent (no matter
+    whether the NaN is in the real or imaginary part). As the representant for
+    the returned array the smallest one in the lexicographical order is chosen.
+    For multi-dimensional inputs, `unique_inverse` is reshaped such that the
+    input can be reconstructed using
+    ``dpnp.take(unique, unique_inverse, axis=axis)``.
 
     Examples
     --------
@@ -2272,7 +2278,6 @@ def unique(
     """
 
     if axis is None:
-        ar = dpnp.ravel(ar)
         return _unique_1d(
             ar, return_index, return_inverse, return_counts, equal_nan
         )

tests/test_manipulation.py

@@ -629,6 +629,17 @@ def test_2d_axis(self, dt, axis_kwd, return_kwds):
         if len(return_kwds) == 0:
             assert_array_equal(result, expected)
         else:
+            if (
+                len(axis_kwd) == 0
+                and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1"
+            ):
+                # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+                # returned flatten unique_inverse till 2.0.1 version
+                expected = (
+                    expected[:2]
+                    + (expected[2].reshape(a.shape),)
+                    + expected[3:]
+                )
             for iv, v in zip(result, expected):
                 assert_array_equal(iv, v)
 
@@ -665,6 +676,11 @@ def test_2d_axis_inverse(self, axis):
 
         result = dpnp.unique(ia, return_inverse=True, axis=axis)
         expected = numpy.unique(a, return_inverse=True, axis=axis)
+        if axis is None and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1":
+            # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+            # returned flatten unique_inverse till 2.0.1 version
+            expected = expected[:1] + (expected[1].reshape(a.shape),)
+
         for iv, v in zip(result, expected):
             assert_array_equal(iv, v)
 
@@ -776,5 +792,16 @@ def test_2d_axis_nans(self, dt, axis_kwd, return_kwds, row):
         if len(return_kwds) == 0:
             assert_array_equal(result, expected)
         else:
+            if (
+                len(axis_kwd) == 0
+                and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1"
+            ):
+                # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+                # returned flatten unique_inverse till 2.0.1 version
+                expected = (
+                    expected[:2]
+                    + (expected[2].reshape(a.shape),)
+                    + expected[3:]
+                )
             for iv, v in zip(result, expected):
                 assert_array_equal(iv, v)

tests/test_sycl_queue.py

@@ -2408,6 +2408,10 @@ def test_unique(axis, device):
 
     result = dpnp.unique(ia, True, True, True, axis=axis)
     expected = numpy.unique(a, True, True, True, axis=axis)
+    if axis is None and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1":
+        # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+        # returned flatten unique_inverse till 2.0.1 version
+        expected = expected[:2] + (expected[2].reshape(a.shape),) + expected[3:]
     for iv, v in zip(result, expected):
         assert_array_equal(iv, v)
 

tests/third_party/cupy/manipulation_tests/test_add_remove.py

@@ -184,7 +184,12 @@ def test_unique_index(self, xp, dtype):
     @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]
+        result = xp.unique(a, return_inverse=True)[1]
+        if xp is numpy and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1":
+            # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+            # returned flatten unique_inverse till 2.0.1 version
+            result = result.reshape(a.shape)
+        return result
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
@@ -208,9 +213,14 @@ def test_unique_counts(self, xp, dtype):
     @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(
+        result = xp.unique(
             a, return_index=True, return_inverse=True, return_counts=True
         )
+        if xp is numpy and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1":
+            # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+            # returned flatten unique_inverse till 2.0.1 version
+            result = result[:2] + (result[2].reshape(a.shape),) + result[3:]
+        return result
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()
@@ -240,9 +250,14 @@ def test_unique_empty(self, xp, dtype):
     @testing.numpy_cupy_array_equal()
     def test_unique_empty_return_all_no_axis(self, xp, dtype):
         a = xp.empty((3, 0, 2), dtype=dtype)
-        return xp.unique(
+        result = xp.unique(
             a, return_index=True, return_inverse=True, return_counts=True
         )
+        if xp is numpy and numpy.lib.NumpyVersion(numpy.__version__) < "2.0.1":
+            # gh-26961: numpy.unique(..., return_inverse=True, axis=None)
+            # returned flatten unique_inverse till 2.0.1 version
+            result = result[:2] + (result[2].reshape(a.shape),) + result[3:]
+        return result
 
     @testing.for_all_dtypes(no_float16=True, no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()