anf_error.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, relay.If(cond, net, y))
net = relay.add(net, z)
net = relay.Function([cond,x,y,z], net)
net = relay.ir_pass.infer_type(net)

mod = relay.Module.from_expr(net)
mod = relay.transform.ToANormalForm()(mod)
mod[mod.entry_func] = relay.ir_pass.gradient(mod[mod.entry_func], mode='higher_order')
print(mod[mod.entry_func])

mod = relay.transform.ToANormalForm()(mod) # error

mod[mod.entry_func] = relay.transform.ToCPS(mod[mod.entry_func])
mod = relay.transform.PartialEvaluate()(mod)
mod = relay.transform.DeadCodeElimination()(mod)
print(mod[mod.entry_func])

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')

with relay.build_config(opt_level=3):
    aut = relay.create_executor('vm', mod=mod).evaluate(mod[mod.entry_func])(cond_np, x_np, y_np, z_np)