Performs sorting along the given axis and returns an array of indices having same shape as an input array that index data in sorted order. Parameters ---------- data: tvm.te.Tensor The input array. axis : int, optional Axis long which to sort the input tensor.
(data, axis=-1, is_ascend=1, dtype="float32", ret_type="indices", workspace=None)
| 842 | |
| 843 | |
| 844 | def argsort_thrust(data, axis=-1, is_ascend=1, dtype="float32", ret_type="indices", workspace=None): |
| 845 | """Performs sorting along the given axis and returns an array of indices |
| 846 | having same shape as an input array that index data in sorted order. |
| 847 | |
| 848 | Parameters |
| 849 | ---------- |
| 850 | data: tvm.te.Tensor |
| 851 | The input array. |
| 852 | |
| 853 | axis : int, optional |
| 854 | Axis long which to sort the input tensor. |
| 855 | |
| 856 | is_ascend : boolean, optional |
| 857 | Whether to sort in ascending or descending order. |
| 858 | |
| 859 | dtype : string, optional |
| 860 | DType of the output indices. |
| 861 | |
| 862 | ret_type : string, optional |
| 863 | The return type [both, indices]. |
| 864 | "both": return both sorted data and indices. |
| 865 | "indices": return sorted indices only. |
| 866 | |
| 867 | workspace : Optional[tvm.te.Tensor] |
| 868 | A buffer to store intermediate results. The size of the workspace should be sufficiently |
| 869 | large, this can be obtained by overestimation or memory usage profiling. If None, it will |
| 870 | fallback to use thrust internal memory allocation. |
| 871 | |
| 872 | Returns |
| 873 | ------- |
| 874 | out : tvm.te.Tensor |
| 875 | The output of this function. |
| 876 | """ |
| 877 | return topk_thrust(data, 0, axis, ret_type, is_ascend, dtype, workspace) |
| 878 | |
| 879 | |
| 880 | def topk(data, k=1, axis=-1, ret_type="both", is_ascend=False, dtype="int64"): |
nothing calls this directly
no test coverage detected
searching dependent graphs…