diff options
| author | Matthew <quitesimplymatt@gmail.com> | 2022-01-11 17:04:39 +0000 |
|---|---|---|
| committer | Matthew <quitesimplymatt@gmail.com> | 2022-01-13 10:04:23 +0000 |
| commit | 8b967ff2e70afe3a1fd32e33b36f66f34c259139 (patch) | |
| tree | 1f9e76d46214c86a5c6cd328e38cb31d733354b9 /numpy/array_api/_set_functions.py | |
| parent | 813a0c11186ded0b5caeb853fd2b22fb9addd511 (diff) | |
| download | numpy-8b967ff2e70afe3a1fd32e33b36f66f34c259139.tar.gz | |
BUG: Return correctly shaped inverse indices in `array_api`
Specifically for `xp.unique_all()` and `xp.unique_inverse()`
Diffstat (limited to 'numpy/array_api/_set_functions.py')
| -rw-r--r-- | numpy/array_api/_set_functions.py | 20 |
1 files changed, 15 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: |