Function test_simple_add

tests/python/relax/test_binding_rewrite.py:53–75

()

def test_simple_add():
    root_fn = Identity["main"]
    dfb = root_fn.body.blocks[0]

    rwt = DataflowBlockRewrite(dfb, root_fn)
    rwt.add(name="tmp", expr=Identity["main"].params[0], is_dfvar=True)

    assert_immutability(rwt, dfb, root_fn)

    # check "tmp" added
    assert "tmp" in name_to_binding(rwt.mutated_root_fn())

    @tvm.script.ir_module
    class GroundTruth:
        @R.function
        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
            with R.dataflow():
                lv0 = x
                tmp: R.Tensor((32, 32), "float32") = x
                R.output(lv0)
            return lv0

    tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])


def test_simple_auto_add_var():