Merge pull request numpy#26914 from seberg/partial-revert-unique-inverse

API: Partially revert unique with return_inverse

doc/source/release/2.0.0-notes.rst

@@ -1496,6 +1496,12 @@ the ``unique_inverse`` output is now shaped such that the input can be reconstru
 directly using ``np.take(unique, unique_inverse)`` when ``axis=None``, and
 ``np.take_along_axis(unique, unique_inverse, axis=axis)`` otherwise.
 
+.. note::
+    This change was reverted in 2.0.1 except for ``axis=None``.  The correct
+    reconstruction is always ``np.take(unique, unique_inverse, axis=axis)``.
+    When 2.0.0 needs to be supported, add ``unique_inverse.reshape(-1)``
+    to code.
+
 (`gh-25553 <https://github.com/numpy/numpy/pull/25553>`__,
 `gh-25570 <https://github.com/numpy/numpy/pull/25570>`__)
 

numpy/lib/_arraysetops_impl.py

@@ -228,8 +228,13 @@ def unique(ar, return_index=False, return_inverse=False,
     .. versionchanged: 2.0
         For multi-dimensional inputs, ``unique_inverse`` is reshaped
         such that the input can be reconstructed using
-        ``np.take(unique, unique_inverse)`` when ``axis = None``, and
-        ``np.take_along_axis(unique, unique_inverse, axis=axis)`` otherwise.
+        ``np.take(unique, unique_inverse, axis=axis)``. The result is
+        now not 1-dimensional when ``axis=None``.
+
+        Note that in NumPy 2.0.0 a higher dimensional array was returned also
+        when ``axis`` was not ``None``.  This was reverted, but
+        ``inverse.reshape(-1)`` can be used to ensure compatibility with both
+        versions.
 
     Examples
     --------
@@ -282,7 +287,7 @@ def unique(ar, return_index=False, return_inverse=False,
     ar = np.asanyarray(ar)
     if axis is None:
         ret = _unique1d(ar, return_index, return_inverse, return_counts, 
-                        equal_nan=equal_nan, inverse_shape=ar.shape)
+                        equal_nan=equal_nan, inverse_shape=ar.shape, axis=None)
         return _unpack_tuple(ret)
 
     # axis was specified and not None
@@ -328,13 +333,15 @@ def reshape_uniq(uniq):
 
     output = _unique1d(consolidated, return_index,
                        return_inverse, return_counts,
-                       equal_nan=equal_nan, inverse_shape=inverse_shape)
+                       equal_nan=equal_nan, inverse_shape=inverse_shape,
+                       axis=axis)
     output = (reshape_uniq(output[0]),) + output[1:]
     return _unpack_tuple(output)
 
 
 def _unique1d(ar, return_index=False, return_inverse=False,
-              return_counts=False, *, equal_nan=True, inverse_shape=None):
+              return_counts=False, *, equal_nan=True, inverse_shape=None,
+              axis=None):
     """
     Find the unique elements of an array, ignoring shape.
     """
@@ -371,7 +378,7 @@ def _unique1d(ar, return_index=False, return_inverse=False,
         imask = np.cumsum(mask) - 1
         inv_idx = np.empty(mask.shape, dtype=np.intp)
         inv_idx[perm] = imask
-        ret += (inv_idx.reshape(inverse_shape),)
+        ret += (inv_idx.reshape(inverse_shape) if axis is None else inv_idx,)
     if return_counts:
         idx = np.concatenate(np.nonzero(mask) + ([mask.size],))
         ret += (np.diff(idx),)

numpy/lib/tests/test_arraysetops.py

@@ -828,11 +828,8 @@ def test_unique_1d_with_axis(self, axis):
     def test_unique_inverse_with_axis(self, axis):
         x = np.array([[4, 4, 3], [2, 2, 1], [2, 2, 1], [4, 4, 3]])
         uniq, inv = unique(x, return_inverse=True, axis=axis)
-        assert_equal(inv.ndim, x.ndim)
-        if axis is None:
-            assert_array_equal(x, np.take(uniq, inv))
-        else:
-            assert_array_equal(x, np.take_along_axis(uniq, inv, axis=axis))
+        assert_equal(inv.ndim, x.ndim if axis is None else 1)
+        assert_array_equal(x, np.take(uniq, inv, axis=axis))
 
     def test_unique_axis_zeros(self):
         # issue 15559
@@ -844,7 +841,7 @@ def test_unique_axis_zeros(self):
         assert_equal(uniq.dtype, single_zero.dtype)
         assert_array_equal(uniq, np.empty(shape=(1, 0)))
         assert_array_equal(idx, np.array([0]))
-        assert_array_equal(inv, np.array([[0], [0]]))
+        assert_array_equal(inv, np.array([0, 0]))
         assert_array_equal(cnt, np.array([2]))
 
         # there's 0 elements of shape (2,) along axis 1
@@ -854,7 +851,7 @@ def test_unique_axis_zeros(self):
         assert_equal(uniq.dtype, single_zero.dtype)
         assert_array_equal(uniq, np.empty(shape=(2, 0)))
         assert_array_equal(idx, np.array([]))
-        assert_array_equal(inv, np.empty((1, 0)))
+        assert_array_equal(inv, np.array([]))
         assert_array_equal(cnt, np.array([]))
 
         # test a "complicated" shape
@@ -923,7 +920,7 @@ def _run_axis_tests(self, dtype):
         msg = "Unique's return_index=True failed with axis=0"
         assert_array_equal(data[idx], uniq, msg)
         msg = "Unique's return_inverse=True failed with axis=0"
-        assert_array_equal(np.take_along_axis(uniq, inv, axis=0), data)
+        assert_array_equal(np.take(uniq, inv, axis=0), data)
         msg = "Unique's return_counts=True failed with axis=0"
         assert_array_equal(cnt, np.array([2, 2]), msg)
 
@@ -932,7 +929,7 @@ def _run_axis_tests(self, dtype):
         msg = "Unique's return_index=True failed with axis=1"
         assert_array_equal(data[:, idx], uniq)
         msg = "Unique's return_inverse=True failed with axis=1"
-        assert_array_equal(np.take_along_axis(uniq, inv, axis=1), data)
+        assert_array_equal(np.take(uniq, inv, axis=1), data)
         msg = "Unique's return_counts=True failed with axis=1"
         assert_array_equal(cnt, np.array([2, 1, 1]), msg)