code: def test_after_partial_eval():
      """Test transformation following reverse mode ad and PartialEval"""
      mod = tvm.IRModule()
  
      shape = (10, 10)
      dtype = "float32"
      t = relay.TensorType(shape, dtype)
  
      x = relay.var("x", t)
      y = relay.var("y", t)
  
      func = relay.Function(
          [x, y], (x * y) * relay.const(np.ones(shape, dtype))
      )
      func = run_infer_type(func)
      back_func = transform.gradient(func)
      back_func = run_infer_type(back_func)
  
      mod["main"] = back_func
      back_func = mod["main"]
  
      seq = tvm.transform.Sequential(
          [
              transform.PartialEvaluate(),
              transform.InferType(),
              transform.LazyGradientInit(),
              transform.InferType(),
              transform.DeadCodeElimination(),
              transform.InferType(),
          ]
      )
  
      mod = seq(mod)
  
      assert mod["main"].checked_type == relay.FuncType(
          [t, t], relay.TupleType([t, relay.TupleType([t, t])])
      )
  
      x = rand(dtype, *shape)
      y = rand(dtype, *shape)
      (forward), (
          grad_x,
          grad_y,
      ) = create_executor(mod=mod).evaluate(
          back_func
      )(x, y)
      assert_allclose(forward.numpy(), x.numpy() * y.numpy())
      assert_allclose(grad_x.numpy(), y.numpy())
      assert_allclose(grad_y.numpy(), x.numpy())
  
imports: import numpy
  import tvm
  from tvm import relay
  from tvm.relay import create_executor, transform
  from tvm.relay.testing import rand, run_infer_type
  import tvm.testing
  from tvm.testing import assert_allclose
label_desc: Use image node_id 10 for calling a global function with example usage: test_after_partial_eval() without return types