diff --git a/backends/arm/_passes/fuse_identical_input_transforms_pass.py b/backends/arm/_passes/fuse_identical_input_transforms_pass.py index 39bc807ce0c..00e6f5d3736 100644 --- a/backends/arm/_passes/fuse_identical_input_transforms_pass.py +++ b/backends/arm/_passes/fuse_identical_input_transforms_pass.py @@ -25,3 +25,4 @@ class FuseIdenticalInputTransformsPass( # Both bases declare this; restate it once so the two agree. target_ops: set[Any] = _FuseIdenticalInputTransformsPass.target_ops + _recompile_before_retrace = False diff --git a/backends/transforms/fuse_identical_input_transforms_pass.py b/backends/transforms/fuse_identical_input_transforms_pass.py index cc8da817cce..b74e35c9931 100644 --- a/backends/transforms/fuse_identical_input_transforms_pass.py +++ b/backends/transforms/fuse_identical_input_transforms_pass.py @@ -152,6 +152,7 @@ class FuseIdenticalInputTransformsPass(ExportPass): _ELEMENTWISE_OPS = _BINARY_ELEMENTWISE_OPS | _NARY_ELEMENTWISE_OPS target_ops = _ELEMENTWISE_OPS | _CONCAT_OPS + _recompile_before_retrace = True def __init__(self, exported_program: ExportedProgram | None = None) -> None: super().__init__() @@ -179,7 +180,8 @@ def call(self, graph_module: GraphModule) -> PassResult: if modified: graph_module.graph.eliminate_dead_code() graph_module.graph.lint() - graph_module.recompile() + if self._recompile_before_retrace: + graph_module.recompile() graph_module = super().call(graph_module).graph_module return PassResult(graph_module, modified)