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 or tvm.te.Tensor, 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"]
      data_buf = tvm.tir.decl_buffer(
          data.shape, data.dtype, "data_buf", data_alignment=8
      )
      out_shape = list(get_const_tuple(data.shape))
      kvar = tvm.te.size_var("k")
      if not isinstance(k, int):
          out_shape[axis] = kvar
      elif k >= 1:
          out_shape[axis] = k
      out_bufs = []
      if ret_type in ["both", "values"]:
          out_bufs.append(
              tvm.tir.decl_buffer(
                  out_shape, data.dtype, "value_buf", data_alignment=8
              )
          )
      if ret_type in ["both", "indices"]:
          out_bufs.append(
              tvm.tir.decl_buffer(
                  out_shape, dtype, "indices_buf", data_alignment=8
              )
          )
      out_shapes = [out_shape] * len(out_bufs)
  
      kv = kvar if not isinstance(k, int) else k
      out = te.extern(
          out_shapes,
          [data],
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.sort.topk",
              ins[0],
              *outs,
              kv,
              axis,
              ret_type,
              is_ascend,
          ),
          in_buffers=[data_buf],
          out_buffers=out_bufs,
          name="topk_cpu",
          tag="topk_cpu",
      )
      return out
  
imports: import tvm
  from tvm import te
  from .utils import get_const_tuple
label_desc: Use image node_id 3 for calling a global function with example usage: topk(data, k, axis, ret_type, is_ascend, dtype) and returns: out