コード例 #1
0
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")
コード例 #2
0
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")