code: def test_fold_const_with_on_device():
      """Make sure on_device annotations don't get in the way of constant folding"""
      c_data = np.array([1, 2, 3]).astype("float32")
      t = relay.TensorType([1, 2, 3], "float32")
  
      def before():
          c = relay.const(c_data)
          x = relay.var("x", t)
          x.virtual_device_ = tvm.cpu()
          y = relay.add(c, c)
          y = relay.multiply(y, relay.const(2, "float32"))
          y = relay.add(x, y)
          z = relay.add(y, c)
          f = relay.Function([x], z)
          f.virtual_device_ = tvm.cpu()
          return f
  
      def expected():
          x = relay.var("x", t)
          x.virtual_device_ = tvm.cpu()
          c_folded = (c_data + c_data) * 2
          y = relay.add(x, relay.const(c_folded))
          z = relay.add(y, relay.const(c_data))
          f = relay.Function([x], z)
          f.virtual_device_ = tvm.cpu()
          return f
  
      zz = run_opt_pass(before(), transform.FoldConstant())
      zexpected = run_opt_pass(expected(), transform.InferType())
      tvm.ir.assert_structural_equal(zz, zexpected)
  
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import relay
  from tvm.relay.backend import Executor
  from tvm.relay import transform
  from tvm.relay.build_module import bind_params_by_name
  from tvm.relay.testing import run_infer_type, create_workload
label_desc: Use image node_id 5 for calling a global function with example usage: test_fold_const_with_on_device() and returns: f, f