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")
| 769 | |
| 770 | |
| 771 | def argsort(data, axis=-1, is_ascend=1, dtype="float32", ret_type="indices"): |
| 772 | """Performs sorting along the given axis and returns an array of indices |
| 773 | having same shape as an input array that index data in sorted order. |
| 774 | |
| 775 | Parameters |
| 776 | ---------- |
| 777 | data: tvm.te.Tensor |
| 778 | The input array. |
| 779 | |
| 780 | axis : int, optional |
| 781 | Axis long which to sort the input tensor. |
| 782 | |
| 783 | is_ascend : boolean, optional |
| 784 | Whether to sort in ascending or descending order. |
| 785 | |
| 786 | dtype : string, optional |
| 787 | DType of the output indices. |
| 788 | |
| 789 | ret_type : string, optional |
| 790 | The return type [both, indices]. |
| 791 | "both": return both sorted data and indices. |
| 792 | "indices": return sorted indices only. |
| 793 | |
| 794 | Returns |
| 795 | ------- |
| 796 | out : tvm.te.Tensor |
| 797 | The output of this function. |
| 798 | """ |
| 799 | ndim = len(data.shape) |
| 800 | axis = ndim + axis if axis < 0 else axis |
| 801 | if axis != ndim - 1: |
| 802 | # Prepare for sorting along axis -1. |
| 803 | axes = swap(list(range(ndim)), axis) |
| 804 | data = transpose(data, axes) |
| 805 | |
| 806 | value_buf = tvm.tirx.decl_buffer( |
| 807 | data.shape, data.dtype, "value_buf", data_alignment=8, layout=None |
| 808 | ) |
| 809 | value_swap_buf = tvm.tirx.decl_buffer( |
| 810 | data.shape, data.dtype, "value_swap_buf", data_alignment=8, layout=None |
| 811 | ) |
| 812 | indices_buf = tvm.tirx.decl_buffer(data.shape, dtype, "out_buf", data_alignment=8, layout=None) |
| 813 | indices_swap_buf = tvm.tirx.decl_buffer( |
| 814 | data.shape, dtype, "out_swap_buf", data_alignment=8, layout=None |
| 815 | ) |
| 816 | |
| 817 | outs = te.extern( |
| 818 | [data.shape, data.shape, data.shape, data.shape], |
| 819 | [data], |
| 820 | lambda ins, outs: sort_ir( |
| 821 | ins[0], |
| 822 | outs[0], |
| 823 | outs[2], |
| 824 | -1, |
| 825 | is_ascend, |
| 826 | indices_out=outs[1], |
| 827 | indices_out_swap=outs[3], |
| 828 | ), |
no test coverage detected
searching dependent graphs…