code: def sort_thrust(data, axis=-1, is_ascend=1):
      """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, optional
          Whether to sort in ascending or descending order.
  
      Returns
      -------
      out : tvm.te.Tensor
          The output of this function.
      """
      dtype = "float32"
  
      ndim = len(data.shape)
      axis = ndim + axis if axis < 0 else axis
  
      if axis != ndim - 1:
          # Prepare for sorting along axis -1.
          axes = swap(list(range(ndim)), axis)
          data = transpose(data, axes)
  
      value_buf = tvm.tir.decl_buffer(
          data.shape, data.dtype, "value_buf", data_alignment=8
      )
      indices_buf = tvm.tir.decl_buffer(
          data.shape, dtype, "out_buf", data_alignment=8
      )
      out = te.extern(
          [data.shape, data.shape],
          [data],
          ## TODO(mbrookhart): This thrust function is actually doing argsort, not sort
          ## For performance, we should probably rename the contrib function and add
          ## a pure sort
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.thrust.sort",
              ins[0],
              outs[0],
              outs[1],
              is_ascend,
          ),
          out_buffers=[value_buf, indices_buf],
          name="sort_gpu",
          tag="sort_gpu",
      )[0]
  
      if axis != ndim - 1:
          axes = swap(list(range(ndim)), axis)
          out = transpose(out, axes)
      return out
  
imports: import tvm
  from tvm import te
  from .injective import schedule_injective_from_existing
  from ..transform import strided_slice, transpose
  from ..None import tag
  from ..utils import ceil_div, swap
  from ..math import cast, ceil_log2
label_desc: Use image node_id 9 for calling a global function with example usage: sort_thrust(data, axis, is_ascend) and returns: out