sort inp in the specfic axis, and return the sorted value and index for example: inp: [[ 1.7783 -1.8184 1.0701] [-0.0712 -1.4623 1.3243]] axis: -1 descend: True return: [[ 1.7783 1.0701 -1.8184] [ 1.3243 -0.0712 -1.4623]] [[0 2 1]
(inp, axis=-1, descending=True, is_stable=True)
| 327 | |
| 328 | |
| 329 | def argsort(inp, axis=-1, descending=True, is_stable=True): |
| 330 | """ |
| 331 | sort inp in the specfic axis, and return the sorted value and index |
| 332 | for example: |
| 333 | inp: |
| 334 | [[ 1.7783 -1.8184 1.0701] |
| 335 | [-0.0712 -1.4623 1.3243]] |
| 336 | axis: -1 |
| 337 | descend: True |
| 338 | return: |
| 339 | [[ 1.7783 1.0701 -1.8184] |
| 340 | [ 1.3243 -0.0712 -1.4623]] |
| 341 | [[0 2 1] |
| 342 | [2 0 1]] |
| 343 | """ |
| 344 | axis = axis + inp.ndim if axis < 0 else axis |
| 345 | idx = iota(np.int32, inp.shape, axis) |
| 346 | return _sort_according_to_key( |
| 347 | inp, idx, axis=axis, descending=descending, is_stable=is_stable |
| 348 | ) |
| 349 | |
| 350 | |
| 351 | @register_lower_rule(mops.Argsort) |
no test coverage detected