Skip to content

Instantly share code, notes, and snippets.

@vinx13
Last active June 28, 2019 16:59
Show Gist options
  • Save vinx13/ec07cf196a1e48e4aa9eaca684feb9b6 to your computer and use it in GitHub Desktop.
Save vinx13/ec07cf196a1e48e4aa9eaca684feb9b6 to your computer and use it in GitHub Desktop.
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment