code: def test_trace_expr_sum_custom():
      @tvm.register_func("tvm.tir.trace_callback4")
      def trace_buffer(x):
          return
  
      def check_expr_sum_custom(dtype):
          n = 4
          a = te.placeholder((n, n), name="a", dtype=dtype)
          b = te.placeholder((n, n), name="b", dtype=dtype)
          c = te.compute(
              a.shape,
              lambda i, j: tvm.tir.trace(
                  [a[i][j]], "tvm.tir.trace_callback4"
              )
              + tvm.tir.trace([b[i][j]], "tvm.tir.trace_callback4"),
          )
          s = te.create_schedule(c.op)
          f = tvm.build(s, [a, b, c])
          npa = np.array(
              [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
              dtype=a.dtype,
          )
          npb = np.array(
              [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
              dtype=a.dtype,
          )
          xnd = tvm.nd.array(npa)
          ynd = tvm.nd.array(npb)
          znd = tvm.nd.array(np.zeros((n, n), dtype=c.dtype))
          f(xnd, ynd, znd)
          assert np.array_equal(znd.numpy(), npa + npb)
  
      for t in ["float64", "float32", "int64", "int32"]:
          check_expr_sum_custom(t)
  
file_path: tvm/tests/python/runtime/test_runtime_trace.py
imports: import tvm
  from tvm import te
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: test_trace_expr_sum_custom() without return types