def check_result(args, expected_result, mod=None):
"""
Check that evaluating `expr` applied to the arguments produces
`result` on Relay VM.
Parameters
----------
args: list of Expr
The arguments to supply the expr.
expected_result:
The expected result of running the expression.
"""
for target, dev in tvm.testing.enabled_targets():
vm = relay.create_executor("vm", device=dev, target=target, mod=mod)
rts_result = vm.evaluate()(*args)
tvm.testing.assert_allclose(expected_result, rts_result.asnumpy())
def test_multiple_ifs():
mod = tvm.IRModule({})
b = relay.var("b")
v0 = relay.var("v0")
v1 = relay.var("v1")
v2 = relay.var("v2")
v3 = relay.var("v3")
out = relay.Tuple([v2, v3])
out = relay.Let(v3, relay.If(b, v1, v0), out)
out = relay.Let(v2, relay.If(b, v0, v1), out)
out = relay.Let(v1, relay.Tuple([relay.const(1)]), out)
out = relay.Let(v0, relay.Tuple([relay.const(0)]), out)
fn = relay.Function([b], out)
mod["main"] = fn
dev = tvm.runtime.device("llvm", 0)
vm = relay.create_executor(device=dev, mod=mod, kind="vm")
res = vmobj_to_list(vm.evaluate()(False))
assert res == [1, 0]