Function test_variable_replacer

tests/python/tirx-transform/test_tir_functor.py:228–243

Test expression mutator that replaces variables

()

def test_variable_replacer():
    """Test expression mutator that replaces variables"""
    x = Var("x", dtype="int32")
    y = Var("y", dtype="int32")
    expr = Add(x, Mul(y, IntImm("int32", 3)))

    replacer = VariableReplacer({"x": 10, "y": 5})
    result = replacer.visit_expr(expr)

    # Should be Add(IntImm(10), Mul(IntImm(5), IntImm(3)))
    assert isinstance(result, Add)
    assert isinstance(result.a, IntImm)
    assert result.a.value == 10
    assert isinstance(result.b, Mul)
    assert isinstance(result.b.a, IntImm)
    assert result.b.a.value == 5


def test_add_to_sub_mutator():