Performs sorting along the given axis and returns an array of sorted values with the same shape as the input data. Parameters ---------- data: tvm.te.Tensor The input array. axis : int, optional Axis long which to sort the input tensor. is_ascend : boolean,
(data, axis=-1, is_ascend=1)
| 655 | |
| 656 | |
| 657 | def sort(data, axis=-1, is_ascend=1): |
| 658 | """Performs sorting along the given axis and returns an array of |
| 659 | sorted values with the same shape as the input data. |
| 660 | |
| 661 | Parameters |
| 662 | ---------- |
| 663 | data: tvm.te.Tensor |
| 664 | The input array. |
| 665 | |
| 666 | axis : int, optional |
| 667 | Axis long which to sort the input tensor. |
| 668 | |
| 669 | is_ascend : boolean, optional |
| 670 | Whether to sort in ascending or descending order. |
| 671 | |
| 672 | Returns |
| 673 | ------- |
| 674 | out : tvm.te.Tensor |
| 675 | The output of this function. |
| 676 | """ |
| 677 | ndim = len(data.shape) |
| 678 | axis = ndim + axis if axis < 0 else axis |
| 679 | if axis != ndim - 1: |
| 680 | # Prepare for sorting along axis -1. |
| 681 | axes = swap(list(range(ndim)), axis) |
| 682 | data = transpose(data, axes) |
| 683 | |
| 684 | value_buf = tvm.tirx.decl_buffer( |
| 685 | data.shape, data.dtype, "value_buf", data_alignment=8, layout=None |
| 686 | ) |
| 687 | value_buf_swap = tvm.tirx.decl_buffer( |
| 688 | data.shape, data.dtype, "value_buf_swap", data_alignment=8, layout=None |
| 689 | ) |
| 690 | |
| 691 | out = te.extern( |
| 692 | [data.shape, data.shape], |
| 693 | [data], |
| 694 | lambda ins, outs: sort_ir(ins[0], outs[0], outs[1], -1, is_ascend), |
| 695 | out_buffers=[value_buf, value_buf_swap], |
| 696 | name="sort_gpu", |
| 697 | tag="sort_gpu", |
| 698 | )[0] |
| 699 | |
| 700 | if axis != ndim - 1: |
| 701 | axes = swap(list(range(ndim)), axis) |
| 702 | out = transpose(out, axes) |
| 703 | |
| 704 | return out |
| 705 | |
| 706 | |
| 707 | def sort_thrust(data, axis=-1, is_ascend=1, workspace=None): |