code: def test_sort():
      """Tests sort function"""
      n = 2
      l = 5
      m = 3
      data = te.placeholder((n, l, m), name="data")
      sort_num = te.placeholder((n, m), name="sort_num", dtype="int32")
      axis = 1
      is_ascend = False
      out = te.extern(
          data.shape,
          [data, sort_num],
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.sort.argsort_nms",
              ins[0],
              ins[1],
              outs[0],
              axis,
              is_ascend,
          ),
          dtype="int32",
          name="sort_tensor",
      )
      input_data = [
          [
              [1, 2, 3],
              [2, 4.5, 3.5],
              [1.1, 0.5, 1],
              [3.2, -5, 0.5],
              [1.5, 0, 0],
          ],
          [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]],
      ]
      sort_num_input = [[1, 2, 3], [4, 5, 5]]
      sorted_index = [
          [[0, 1, 1], [1, 0, 0], [2, 2, 2], [3, 3, 3], [4, 4, 4]],
          [[3, 4, 4], [2, 3, 3], [1, 2, 2], [0, 1, 1], [4, 0, 0]],
      ]
  
      dev = tvm.cpu(0)
      target = "llvm"
      s = te.create_schedule(out.op)
      f = tvm.build(s, [data, sort_num, out], target)
      a = tvm.nd.array(np.array(input_data).astype(data.dtype), dev)
      b = tvm.nd.array(
          np.array(sort_num_input).astype(sort_num.dtype), dev
      )
      c = tvm.nd.array(np.zeros(a.shape, dtype=out.dtype), dev)
      f(a, b, c)
      tvm.testing.assert_allclose(
          c.numpy(), np.array(sorted_index).astype(out.dtype), rtol=1e-5
      )
  
file_path: tvm/tests/python/contrib/test_sort.py
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.topi.cuda import sort_by_key
label_desc: Use image node_id 1 for calling a global function with example usage: test_sort() without return types