r"""Returns the indices that would sort the input tensor. Args: inp: input tensor. If it's 2d, the result would be array of indices show how to sort each row in the input tensor. descending: sort in descending order, where the largest comes first. Default: False inp: Ten
(inp: Tensor, descending: bool = False)
| 578 | |
| 579 | |
| 580 | def argsort(inp: Tensor, descending: bool = False) -> Tensor: |
| 581 | r"""Returns the indices that would sort the input tensor. |
| 582 | |
| 583 | Args: |
| 584 | inp: input tensor. If it's 2d, the result would be array of indices show how to sort each row in the input tensor. |
| 585 | descending: sort in descending order, where the largest comes first. Default: False |
| 586 | inp: Tensor: |
| 587 | descending: bool: |
| 588 | |
| 589 | Returns: |
| 590 | indices of int32 indicates how to sort the input. |
| 591 | |
| 592 | Examples: |
| 593 | >>> import numpy as np |
| 594 | >>> x = Tensor(np.array([1,2], dtype=np.float32)) |
| 595 | >>> F.argsort(x) |
| 596 | Tensor([0 1], dtype=int32, device=xpux:0) |
| 597 | """ |
| 598 | assert len(inp.shape) <= 2, "Input should be 1d or 2d" |
| 599 | if descending: |
| 600 | order = "descending" |
| 601 | else: |
| 602 | order = "ascending" |
| 603 | |
| 604 | op = builtin.Argsort(order=order) |
| 605 | if len(inp.shape) == 1: |
| 606 | inp = inp.reshape(1, -1) |
| 607 | _, result = apply(op, inp) |
| 608 | return result[0] |
| 609 | _, result = apply(op, inp) |
| 610 | return result |
| 611 | |
| 612 | |
| 613 | def sort(inp: Tensor, descending: bool = False, dim: int = -1) -> Tuple[Tensor, Tensor]: |