code: def topk(
      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 = axis + ndim if axis < 0 else axis
      assert 0 <= axis < ndim
      dshape = data.shape
      if axis != ndim - 1:
          axes = swap(list(range(ndim)), axis)
          data = transpose(data, axes)
  
      values_buf = tvm.tir.decl_buffer(
          data.shape, data.dtype, "values_buf", data_alignment=8
      )
      values_swap_buf = tvm.tir.decl_buffer(
          data.shape, data.dtype, "values_swap_buf", data_alignment=8
      )
      indices_buf = tvm.tir.decl_buffer(
          data.shape, dtype, "indices_buf", data_alignment=8
      )
      indices_swap_buf = tvm.tir.decl_buffer(
          data.shape, dtype, "indies_swap_buf", data_alignment=8
      )
  
      if ret_type == "values":
          output = te.extern(
              [data.shape, data.shape],
              [data],
              lambda ins, outs: sort_ir(
                  ins[0], outs[0], outs[1], -1, is_ascend
              ),
              out_buffers=[values_buf, values_swap_buf],
              name="topk_gpu",
              tag="topk_gpu",
          )[0]
          if axis != ndim - 1:
              axes = swap(list(range(ndim)), axis)
              output = transpose(output, axes)
      else:
          output = te.extern(
              [data.shape, data.shape, data.shape, data.shape],
              [data],
              lambda ins, outs: sort_ir(
                  ins[0],
                  outs[0],
                  outs[2],
                  -1,
                  is_ascend,
                  indices_out=outs[1],
                  indices_out_swap=outs[3],
              ),
              out_buffers=[
                  values_buf,
                  indices_buf,
                  values_swap_buf,
                  indices_swap_buf,
              ],
              name="topk_gpu",
              tag="topk_gpu",
          )[0:2]
          if axis != ndim - 1:
              axes = swap(list(range(ndim)), axis)
              output[0] = transpose(output[0], axes)
              output[1] = transpose(output[1], axes)
  
      if isinstance(k, int) and k < 1:
          if ret_type == "indices":
              return output[1]
          return output
      beg = [0] * ndim
      end = []
      strides = [1] * ndim
      for i in range(ndim):
          if i == axis:
              end.append(
                  k if isinstance(k, int) else tvm.te.size_var("dim")
              )
          else:
              end.append(dshape[i])
      if ret_type == "both":
          values_out, indices_out = output
          values_out = strided_slice(values_out, beg, end, strides)
          indices_out = strided_slice(indices_out, beg, end, strides)
          output = [values_out, indices_out]
      elif ret_type == "values":
          output = [strided_slice(output, beg, end, strides)]
      else:  # ret_type == "indices"
          indices_out = output[1]
          output = [strided_slice(indices_out, beg, end, strides)]
      return output
  
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 14 for calling a global function with example usage: topk(data, k, axis, ret_type, is_ascend, dtype) and returns: output, output, output