code: def test_exclusive_scan():
      """Tests function test_exclusive_scan"""
      for target in ["cuda", "rocm"]:
          if not tvm.testing.device_enabled(target):
              print("Skip because %s is not enabled" % target)
              continue
  
          with tvm.target.Target(target + " -libs=thrust") as tgt:
              if not thrust_check_func[target](
                  tgt, "tvm.contrib.thrust.sum_scan"
              ):
                  print("skip because thrust is not enabled...")
                  return
  
              for ishape in [(10,), (10, 10), (10, 10, 10)]:
                  values = te.placeholder(
                      ishape, name="values", dtype="int32"
                  )
  
                  scan, reduction = exclusive_scan(
                      values, return_reduction=True
                  )
                  s = schedule_scan([scan, reduction])
  
                  dev = tvm.device(target, 0)
                  f = tvm.build(s, [values, scan, reduction], target)
  
                  values_np = np.random.randint(
                      0, 10, size=ishape
                  ).astype(np.int32)
                  values_np_out = np.zeros(values_np.shape, np.int32)
  
                  if len(ishape) == 1:
                      reduction_shape = ()
                  else:
                      reduction_shape = ishape[:-1]
  
                  reduction_np_out = np.zeros(reduction_shape, np.int32)
  
                  values_in = tvm.nd.array(values_np, dev)
                  values_out = tvm.nd.array(values_np_out, dev)
                  reduction_out = tvm.nd.array(reduction_np_out, dev)
                  f(values_in, values_out, reduction_out)
  
                  ref_values_out = (
                      np.cumsum(values_np, axis=-1, dtype="int32")
                      - values_np
                  )
                  tvm.testing.assert_allclose(
                      values_out.numpy(), ref_values_out, rtol=1e-5
                  )
                  ref_reduction_out = np.sum(values_np, axis=-1)
                  tvm.testing.assert_allclose(
                      reduction_out.numpy(),
                      ref_reduction_out,
                      rtol=1e-5,
                  )
  
file_path: tvm/tests/python/contrib/test_thrust.py
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.topi.cuda import stable_sort_by_key_thrust
  from tvm.topi.cuda.scan import exclusive_scan, scan_thrust, schedule_scan
  from tvm.contrib.thrust import can_use_thrust, can_use_rocthrust
label_desc: Use image node_id 2 for calling a global function with example usage: test_exclusive_scan() without return types