Class VariableReplacer

tests/python/tirx-transform/test_tir_functor.py:107–117

Replace variables with constants

@tirx.functor.mutator
class VariableReplacer(PyStmtExprMutator):
    """Replace variables with constants"""

    def __init__(self, replacements):
        super().__init__()
        self.replacements = replacements

    def visit_var_(self, op: Var):
        if op.name in self.replacements:
            return IntImm("int32", self.replacements[op.name])
        return op


@tirx.functor.mutator