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", workspace=None
)
| 982 | |
| 983 | |
| 984 | def topk_thrust( |
| 985 | data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64", workspace=None |
| 986 | ): |
| 987 | """Get the top k elements in an input tensor along the given axis. |
| 988 | |
| 989 | Parameters |
| 990 | ---------- |
| 991 | data : tvm.te.Tensor |
| 992 | The input tensor. |
| 993 | |
| 994 | k : int, optional |
| 995 | Number of top elements to select. Return all elements if k < 1. |
| 996 | |
| 997 | axis : int, optional |
| 998 | Axis long which to sort the input tensor. |
| 999 | |
| 1000 | ret_type: str, optional |
| 1001 | The return type [both, values, indices]. |
| 1002 | "both": return both top k data and indices. |
| 1003 | "values": return top k data only. |
| 1004 | "indices": return top k indices only. |
| 1005 | |
| 1006 | is_ascend : boolean, optional |
| 1007 | Whether to sort in ascending or descending order. |
| 1008 | |
| 1009 | dtype : string, optional |
| 1010 | The data type of the indices output. |
| 1011 | |
| 1012 | workspace : Optional[tvm.te.Tensor] |
| 1013 | A buffer to store intermediate results. The size of the workspace should be sufficiently |
| 1014 | large, this can be obtained by overestimation or memory usage profiling. If None, it will |
| 1015 | fallback to use thrust internal memory allocation. |
| 1016 | |
| 1017 | Returns |
| 1018 | ------- |
| 1019 | out : tvm.te.Tensor or List[tvm.te.Tensor] |
| 1020 | The computed result. |
| 1021 | """ |
| 1022 | assert ret_type in ["both", "values", "indices"] |
| 1023 | ndim = len(data.shape) |
| 1024 | axis = ndim + axis if axis < 0 else axis |
| 1025 | |
| 1026 | if axis != ndim - 1: |
| 1027 | # Prepare for sorting along axis -1. |
| 1028 | axes = swap(list(range(ndim)), axis) |
| 1029 | data = transpose(data, axes) |
| 1030 | |
| 1031 | data_buf = tvm.tirx.decl_buffer( |
| 1032 | data.shape, data.dtype, "data_buf", data_alignment=8, layout=None |
| 1033 | ) |
| 1034 | if workspace is not None: |
| 1035 | workspace_buf = tvm.tirx.decl_buffer( |
| 1036 | workspace.shape, workspace.dtype, "workspace_buf", data_alignment=8, layout=None |
| 1037 | ) |
| 1038 | else: |
| 1039 | workspace_buf = None |
| 1040 | out_bufs = [ |
| 1041 | tvm.tirx.decl_buffer(data.shape, data.dtype, "value_buf", data_alignment=8, layout=None), |
no test coverage detected
searching dependent graphs…