Get the top k elements in an input tensor along the given axis. Parameters ---------- data : tvm.te.Tensor The input tensor. k : int or tvm.te.Tensor, optional Number of top elements to select. Return all elements if k < 1. axis : int, optional Axis lon
(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64")
| 153 | |
| 154 | |
| 155 | def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"): |
| 156 | """Get the top k elements in an input tensor along the given axis. |
| 157 | |
| 158 | Parameters |
| 159 | ---------- |
| 160 | data : tvm.te.Tensor |
| 161 | The input tensor. |
| 162 | |
| 163 | k : int or tvm.te.Tensor, optional |
| 164 | Number of top elements to select. Return all elements if k < 1. |
| 165 | |
| 166 | axis : int, optional |
| 167 | Axis long which to sort the input tensor. |
| 168 | |
| 169 | ret_type: str, optional |
| 170 | The return type [both, values, indices]. |
| 171 | "both": return both top k data and indices. |
| 172 | "values": return top k data only. |
| 173 | "indices": return top k indices only. |
| 174 | |
| 175 | is_ascend : boolean, optional |
| 176 | Whether to sort in ascending or descending order. |
| 177 | |
| 178 | dtype : string, optional |
| 179 | The data type of the indices output. |
| 180 | |
| 181 | Returns |
| 182 | ------- |
| 183 | out : tvm.te.Tensor or List[tvm.te.Tensor] |
| 184 | The computed result. |
| 185 | """ |
| 186 | assert ret_type in ["both", "values", "indices"] |
| 187 | data_buf = tvm.tirx.decl_buffer( |
| 188 | data.shape, data.dtype, "data_buf", data_alignment=8, layout=None |
| 189 | ) |
| 190 | out_shape = list(get_const_tuple(data.shape)) |
| 191 | kvar = tvm.te.size_var("k") |
| 192 | if not isinstance(k, int): |
| 193 | out_shape[axis] = kvar |
| 194 | elif k >= 1: |
| 195 | out_shape[axis] = k |
| 196 | out_bufs = [] |
| 197 | if ret_type in ["both", "values"]: |
| 198 | out_bufs.append( |
| 199 | tvm.tirx.decl_buffer(out_shape, data.dtype, "value_buf", data_alignment=8, layout=None) |
| 200 | ) |
| 201 | if ret_type in ["both", "indices"]: |
| 202 | out_bufs.append( |
| 203 | tvm.tirx.decl_buffer(out_shape, dtype, "indices_buf", data_alignment=8, layout=None) |
| 204 | ) |
| 205 | out_shapes = [out_shape] * len(out_bufs) |
| 206 | |
| 207 | kv = kvar if not isinstance(k, int) else k |
| 208 | out = te.extern( |
| 209 | out_shapes, |
| 210 | [data], |
| 211 | lambda ins, outs: tvm.tirx.call_packed( |
| 212 | "tvm.contrib.sort.topk", ins[0], *outs, kv, axis, ret_type, is_ascend |
nothing calls this directly
no test coverage detected
searching dependent graphs…