2 files changed, 34 insertions, 5 deletions
diff --git a/numpy/array_api/_set_functions.py b/numpy/array_api/_set_functions.py
index 05ee7e555..db9370f84 100644
--- a/numpy/array_api/_set_functions.py
+++ b/numpy/array_api/_set_functions.py
@@ -41,14 +41,21 @@ def unique_all(x: Array, /) -> UniqueAllResult:
 
     See its docstring for more information.
     """
-    res = np.unique(
+    values, indices, inverse_indices, counts = np.unique(
         x._array,
         return_counts=True,
         return_index=True,
         return_inverse=True,
     )
-
-    return UniqueAllResult(*[Array._new(i) for i in res])
+    # np.unique() flattens inverse indices, but they need to share x's shape
+    # See https://github.com/numpy/numpy/issues/20638
+    inverse_indices = inverse_indices.reshape(x.shape)
+    return UniqueAllResult(
+        Array._new(values),
+        Array._new(indices),
+        Array._new(inverse_indices),
+        Array._new(counts),
+    )
 
 
 def unique_counts(x: Array, /) -> UniqueCountsResult:
@@ -68,13 +75,16 @@ def unique_inverse(x: Array, /) -> UniqueInverseResult:
 
     See its docstring for more information.
     """
-    res = np.unique(
+    values, inverse_indices = np.unique(
         x._array,
         return_counts=False,
         return_index=False,
         return_inverse=True,
     )
-    return UniqueInverseResult(*[Array._new(i) for i in res])
+    # np.unique() flattens inverse indices, but they need to share x's shape
+    # See https://github.com/numpy/numpy/issues/20638
+    inverse_indices = inverse_indices.reshape(x.shape)
+    return UniqueInverseResult(Array._new(values), Array._new(inverse_indices))
 
 
 def unique_values(x: Array, /) -> Array:
diff --git a/numpy/array_api/tests/test_set_functions.py b/numpy/array_api/tests/test_set_functions.py
new file mode 100644
index 000000000..b8eb65d43
--- /dev/null
+++ b/numpy/array_api/tests/test_set_functions.py
@@ -0,0 +1,19 @@
+import pytest
+from hypothesis import given
+from hypothesis.extra.array_api import make_strategies_namespace
+
+from numpy import array_api as xp
+
+xps = make_strategies_namespace(xp)
+
+
+@pytest.mark.parametrize("func", [xp.unique_all, xp.unique_inverse])
+@given(xps.arrays(dtype=xps.scalar_dtypes(), shape=xps.array_shapes()))
+def test_inverse_indices_shape(func, x):
+    """
+    Inverse indices share shape of input array
+
+    See https://github.com/numpy/numpy/issues/20638
+    """
+    out = func(x)
+    assert out.inverse_indices.shape == x.shape