From 8b967ff2e70afe3a1fd32e33b36f66f34c259139 Mon Sep 17 00:00:00 2001 From: Matthew Date: Tue, 11 Jan 2022 17:04:39 +0000 Subject: [PATCH 1/2] BUG: Return correctly shaped inverse indices in `array_api` Specifically for `xp.unique_all()` and `xp.unique_inverse()` --- numpy/array_api/_set_functions.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/numpy/array_api/_set_functions.py b/numpy/array_api/_set_functions.py index 05ee7e555838..db9370f84bd2 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: From ccc1091360639de76fab1c1e2a8b31fda81855fb Mon Sep 17 00:00:00 2001 From: Matthew Date: Thu, 13 Jan 2022 10:03:49 +0000 Subject: [PATCH 2/2] Regression test for inverse indices in `array_api` set functions --- numpy/array_api/tests/test_set_functions.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 numpy/array_api/tests/test_set_functions.py 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 000000000000..b8eb65d430dc --- /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