code: def test_affine_grid():
      def verify_affine_grid(num_batch, target_shape):
          dtype = "float32"
          data_shape = (num_batch, 2, 3)
          data = te.placeholder(data_shape, dtype=dtype)
          out = topi.image.affine_grid(data, target_shape)
  
          @memoize("topi.tests.test_affine_grid.verify_affine_grid")
          def get_ref_data():
              data_np = np.random.uniform(size=data_shape).astype(dtype)
              out_np = tvm.topi.testing.affine_grid_python(
                  data_np, target_shape
              )
              return data_np, out_np
  
          data_np, out_np = get_ref_data()
  
          def check_target(target, dev):
              print("Running on target: %s" % target)
              with tvm.target.Target(target):
                  s = tvm.topi.testing.get_injective_schedule(target)(
                      out
                  )
              tvm_data = tvm.nd.array(data_np, dev)
              tvm_out = tvm.nd.empty(out_np.shape, dtype, dev)
              f = tvm.build(s, [data, out], target)
              f(tvm_data, tvm_out)
  
              tvm.testing.assert_allclose(
                  tvm_out.numpy(), out_np, rtol=1e-5, atol=1e-5
              )
  
          for target, dev in tvm.testing.enabled_targets():
              check_target(target, dev)
  
      verify_affine_grid(1, (16, 32))
      verify_affine_grid(4, (16, 32))
  
file_path: tvm/tests/python/topi/test_topi_image.py
imports: import numpy
  import tvm
  from tvm import te
  from tvm import topi
  import tvm.testing
  import tvm.topi.testing
  from tvm.contrib.pickle_memoize import memoize
label_desc: Use image node_id 6 for calling a global function with example usage: test_affine_grid() and returns: data_np,  out_np