r"""Selects the ``Top-K`` (by default) smallest elements of 2d matrix by row. Args: inp: input tensor. If input tensor is 2d, each row will be sorted. k: number of elements needed. descending: if True, return the largest elements. Default: True kth_only: if True,
(
inp: Tensor,
k: int,
descending: bool = True,
kth_only: bool = False,
no_sort: bool = False,
)
| 653 | |
| 654 | |
| 655 | def topk( |
| 656 | inp: Tensor, |
| 657 | k: int, |
| 658 | descending: bool = True, |
| 659 | kth_only: bool = False, |
| 660 | no_sort: bool = False, |
| 661 | ) -> Tuple[Tensor, Tensor]: |
| 662 | r"""Selects the ``Top-K`` (by default) smallest elements of 2d matrix by row. |
| 663 | |
| 664 | Args: |
| 665 | inp: input tensor. If input tensor is 2d, each row will be sorted. |
| 666 | k: number of elements needed. |
| 667 | descending: if True, return the largest elements. Default: True |
| 668 | kth_only: if True, only the k-th element will be returned. Default: False |
| 669 | no_sort: if True, the returned elements can be unordered. Default: False |
| 670 | |
| 671 | Returns: |
| 672 | tuple of two tensors ``(topk_tensor, indices_of_int32)`` |
| 673 | |
| 674 | Examples: |
| 675 | >>> import numpy as np |
| 676 | >>> x = Tensor(np.array([2, 4, 6, 8, 7, 5, 3, 1], dtype=np.float32)) |
| 677 | >>> top, indices = F.topk(x, 5, descending=False) |
| 678 | >>> print(top.numpy(), indices.numpy()) |
| 679 | [1. 2. 3. 4. 5.] [7 0 6 1 5] |
| 680 | """ |
| 681 | if descending: |
| 682 | k = -k |
| 683 | |
| 684 | if kth_only: |
| 685 | mode = "kth_only" |
| 686 | elif no_sort: |
| 687 | mode = "value_idx_nosort" |
| 688 | else: |
| 689 | mode = "value_idx_sorted" |
| 690 | op = builtin.TopK(mode=mode) |
| 691 | |
| 692 | if not isinstance(k, Tensor): |
| 693 | k = Const(k, "int32", inp.device) |
| 694 | |
| 695 | if len(inp.shape) == 1: |
| 696 | if kth_only: |
| 697 | (tns,) = apply(op, expand_dims(inp, 0), k) |
| 698 | # FIXME: |
| 699 | # could use a dedicated kernel |
| 700 | # gradient may be routed to other indices if k-th value is not unique |
| 701 | ind = argmax((tns == inp).astype("int8")) |
| 702 | tns = squeeze(tns, 0) |
| 703 | else: |
| 704 | tns, ind = apply(op, expand_dims(inp, 0), k) |
| 705 | tns = squeeze(tns, 0) |
| 706 | ind = squeeze(ind, 0) |
| 707 | else: |
| 708 | if kth_only: |
| 709 | (tns,) = apply(op, inp, k) |
| 710 | # FIXME: same as above |
| 711 | ind = argmax((expand_dims(tns, 1) == inp).astype("int8"), 1) |
| 712 | else: |