pe.py

import tvm
import tvm.relay as relay
import tvm.relay.testing
import numpy as np

x = relay.var("x", shape=(1, 16))
y = relay.var("y", shape=(1, 16))
z = relay.var("z", shape=(1, 16))
cond = relay.var("cond", shape=(), dtype='uint1')
net = relay.If(cond, x, y)
net = relay.add(net, z)
net = relay.Function([cond,x,y,z], net)
net = relay.ir_pass.infer_type(net)


cond_np = np.ndarray(()).astype('bool')
x_np = np.random.rand(1, 16).astype('float32')
y_np = np.random.rand(1, 16).astype('float32')
z_np = np.random.rand(1, 16).astype('float32')

back_func_hi = relay.ir_pass.gradient(net, mode='higher_order')
back_func_hi = relay.ir_pass.infer_type(back_func_hi)
back_func_hi = relay.ir_pass.to_a_normal_form(back_func_hi)
back_func_hi = relay.ir_pass.partial_evaluate(back_func_hi)
back_func_hi = relay.ir_pass.infer_type(back_func_hi)
print(back_func_hi)
m = relay.Module.from_expr(back_func_hi)
with relay.build_config(opt_level=3):
    out = relay.create_executor('vm', mod=m).evaluate(back_func_hi)(cond_np, x_np, y_np, z_np)