Get the top k elements in an input tensor along the given axis. Parameters ---------- data : tvm.te.Tensor The input tensor. k : int, optional Number of top elements to select. Return all elements if k < 1. axis : int, optional Axis long which to sort t
(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64")
| 878 | |
| 879 | |
| 880 | def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"): |
| 881 | """Get the top k elements in an input tensor along the given axis. |
| 882 | |
| 883 | Parameters |
| 884 | ---------- |
| 885 | data : tvm.te.Tensor |
| 886 | The input tensor. |
| 887 | |
| 888 | k : int, optional |
| 889 | Number of top elements to select. Return all elements if k < 1. |
| 890 | |
| 891 | axis : int, optional |
| 892 | Axis long which to sort the input tensor. |
| 893 | |
| 894 | ret_type: str, optional |
| 895 | The return type [both, values, indices]. |
| 896 | "both": return both top k data and indices. |
| 897 | "values": return top k data only. |
| 898 | "indices": return top k indices only. |
| 899 | |
| 900 | is_ascend : boolean, optional |
| 901 | Whether to sort in ascending or descending order. |
| 902 | |
| 903 | dtype : string, optional |
| 904 | The data type of the indices output. |
| 905 | |
| 906 | Returns |
| 907 | ------- |
| 908 | out : tvm.te.Tensor or List[tvm.te.Tensor] |
| 909 | The computed result. |
| 910 | """ |
| 911 | assert ret_type in ["both", "values", "indices"] |
| 912 | ndim = len(data.shape) |
| 913 | axis = axis + ndim if axis < 0 else axis |
| 914 | assert 0 <= axis < ndim |
| 915 | dshape = data.shape |
| 916 | if axis != ndim - 1: |
| 917 | axes = swap(list(range(ndim)), axis) |
| 918 | data = transpose(data, axes) |
| 919 | |
| 920 | values_buf = tvm.tirx.decl_buffer( |
| 921 | data.shape, data.dtype, "values_buf", data_alignment=8, layout=None |
| 922 | ) |
| 923 | values_swap_buf = tvm.tirx.decl_buffer( |
| 924 | data.shape, data.dtype, "values_swap_buf", data_alignment=8, layout=None |
| 925 | ) |
| 926 | indices_buf = tvm.tirx.decl_buffer( |
| 927 | data.shape, dtype, "indices_buf", data_alignment=8, layout=None |
| 928 | ) |
| 929 | indices_swap_buf = tvm.tirx.decl_buffer( |
| 930 | data.shape, dtype, "indies_swap_buf", data_alignment=8, layout=None |
| 931 | ) |
| 932 | |
| 933 | if ret_type == "values": |
| 934 | output = te.extern( |
| 935 | [data.shape, data.shape], |
| 936 | [data], |
| 937 | lambda ins, outs: sort_ir(ins[0], outs[0], outs[1], -1, is_ascend), |
nothing calls this directly
no test coverage detected
searching dependent graphs…