Method transform_module

python/tvm/relax/transform/fast_math.py:59–65

(self, mod: IRModule, ctx: tvm.transform.PassContext)

    """

    def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
        fast_math_codegen = FastMathCodeGenerator(mod)
        for gv, func in mod.functions_items():
            if isinstance(func, tvm.relax.Function):
                func = fast_math_codegen.visit_expr(func)
                fast_math_codegen.builder_.update_func(gv, func)
        return fast_math_codegen.builder_.get()