def test_shape_func_nested_function(): @tvm.register_func("relay.ext.test2") def relay_ext_test(func): return None data_shape = (relay.Any(), 16) weight_shape = (relay.Any(), 16) dense = relay.nn.dense( relay.var("data", shape=data_shape), relay.var("weight", shape=weight_shape) ) mod = tvm.IRModule.from_expr(dense) patterns = [("test.dense", is_op("nn.dense")(wildcard(), wildcard()))] passes = tvm.transform.Sequential( [ relay.transform.MergeComposite(patterns), relay.transform.AnnotateTarget(["test2"]), relay.transform.PartitionGraph(), ] ) mod = passes(mod) compiler = VMCompiler() compiler.lower(mod, "llvm")
def test_let_bound_constants(): """This tests for an ICHECK failure for ill-formed IR with let-bound constants""" x = relay.var("x", shape=(3,), dtype="int32") y = relay.take(x, relay.const(0)) z = relay.const(1) f = relay.Function([x], relay.stack((z, y), axis=0)) mod = IRModule.from_expr(f) compiler = VMCompiler() compiler.optimize(mod, target="llvm")