diff options
| author | Charles Harris <charlesr.harris@gmail.com> | 2022-01-13 11:33:17 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-01-13 11:33:17 -0700 |
| commit | 7191d9a4773d77205349ac151f84b72c0ffcf848 (patch) | |
| tree | df7eec8662a2f681ac3bcd20053ba0db864af12d /numpy/array_api/_set_functions.py | |
| parent | f5e74cb4ee7641f72e36901043a7142dd0ac4dcf (diff) | |
| parent | ccc1091360639de76fab1c1e2a8b31fda81855fb (diff) | |
| download | numpy-7191d9a4773d77205349ac151f84b72c0ffcf848.tar.gz | |
Merge pull request #20791 from honno/xp-inverse-indices
BUG: Return correctly shaped inverse indices in `array_api` set functions
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: |