Function test_for

tests/python/tvmscript/test_tvmscript_printer_structural_equal.py:129–152

()

def test_for():
    @T.prim_func(s_tir=True)
    def func1():
        for i, j in T.grid(128, 128):
            with T.sblock():
                pass

    @T.prim_func(s_tir=True)
    def func2():
        for i, j, k in T.grid(128, 128, 128):
            with T.sblock():
                pass

    func1 = func1.with_attr("global_symbol", "main")
    func2 = func2.with_attr("global_symbol", "main")

    with pytest.raises(ValueError) as ve:
        assert_structural_equal(func1, func2)
    assert _error_message(ve.value) == _expected_result(
        func1,
        func2,
        AccessPath.root().attr("body").attr("block").attr("body").attr("body").attr("body"),
        AccessPath.root().attr("body").attr("block").attr("body").attr("body").attr("body"),
    )


if __name__ == "__main__":