code: def topk_thrust(
      data,
      k=1,
      axis=-1,
      ret_type="both",
      is_ascend=False,
      dtype="int64",
  ):
      """Get the top k elements in an input tensor along the given axis.
  
      Parameters
      ----------
      data : tvm.te.Tensor
          The input tensor.
  
      k : int, optional
          Number of top elements to select. Return all elements if k < 1.
  
      axis : int, optional
          Axis long which to sort the input tensor.
  
      ret_type: str, optional
          The return type [both, values, indices].
          "both": return both top k data and indices.
          "values": return top k data only.
          "indices": return top k indices only.
  
      is_ascend : boolean, optional
          Whether to sort in ascending or descending order.
  
      dtype : string, optional
          The data type of the indices output.
  
      Returns
      -------
      out : tvm.te.Tensor or List[tvm.te.Tensor]
          The computed result.
      """
      assert ret_type in ["both", "values", "indices"]
      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)
  
      data_buf = tvm.tir.decl_buffer(
          data.shape, data.dtype, "data_buf", data_alignment=8
      )
      out_bufs = [
          tvm.tir.decl_buffer(
              data.shape, data.dtype, "value_buf", data_alignment=8
          ),
          tvm.tir.decl_buffer(
              data.shape, dtype, "indices_buf", data_alignment=8
          ),
      ]
  
      is_ascend = 1 if is_ascend else 0
  
      out = te.extern(
          [data.shape, data.shape],
          [data],
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.thrust.sort",
              ins[0],
              outs[0],
              outs[1],
              is_ascend,
          ),
          in_buffers=[data_buf],
          out_buffers=out_bufs,
          name="topk_gpu",
          tag="topk_gpu",
      )
  
      if isinstance(k, tvm.tir.IntImm):
          k = k.value
  
      if not isinstance(k, int) or k > 0:
          beg = [0] * ndim
          end = data.shape[:-1] + [
              k if isinstance(k, int) else tvm.te.size_var("dim")
          ]
          strides = [1] * ndim
          out = [strided_slice(o, beg, end, strides) for o in out]
  
      if axis != ndim - 1:
          axes = swap(list(range(ndim)), axis)
          out = [transpose(o, axes) for o in out]
  
      if ret_type == "values":
          out = out[0]
      elif ret_type == "indices":
          out = out[1]
  
      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 15 for calling a global function with example usage: topk_thrust(data, k, axis, ret_type, is_ascend, dtype) and returns: out