Keep output specs in sync in ExportedProgramPassBase

exir/pass_base.py

@@ -10,7 +10,7 @@
 import traceback
 from abc import ABC, abstractmethod
 from contextlib import nullcontext
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import (
     Any,
     Callable,
@@ -38,6 +38,7 @@
 from torch._subclasses.fake_tensor import FakeTensor
 from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode
 from torch.export import ExportedProgram
+from torch.export.graph_signature import ConstantArgument
 from torch.fx import traceback as fx_traceback
 from torch.fx.experimental.proxy_tensor import PythonKeyTracer
 from torch.fx.graph import CodeGen
@@ -237,6 +238,50 @@ class ExportedProgramPassResult:
     modified: bool
 
 
+def _sync_output_specs(exported_program: ExportedProgram) -> bool:
+    """Realign ``output_specs`` with the graph's output node.
+
+    A pass that replaces an output node leaves the signature naming the old one,
+    which later stages then disagree with. Rewriting the spec to match is always
+    safe; changing an output between a node and a literal is not, so that raises
+    rather than guessing.
+
+    Returns whether any spec changed.
+    """
+    output_node = exported_program.graph_module.graph.output_node()
+    outputs = output_node.args[0]
+    assert isinstance(outputs, (tuple, list))
+
+    output_specs = exported_program.graph_signature.output_specs
+    if len(outputs) != len(output_specs):
+        raise ExportPassBaseError(
+            f"Graph has {len(outputs)} outputs, but its signature has "
+            f"{len(output_specs)} output specs"
+        )
+
+    modified = False
+    for output, output_spec in zip(outputs, output_specs):
+        if isinstance(output_spec.arg, ConstantArgument):
+            if isinstance(output, torch.fx.Node):
+                raise ExportPassBaseError(
+                    f"Output {output.name} replaced a literal output; changing output "
+                    "representation is not supported"
+                )
+            if output_spec.arg.value != output:
+                output_spec.arg = replace(output_spec.arg, value=output)
+                modified = True
+            continue
+        if not isinstance(output, torch.fx.Node):
+            raise ExportPassBaseError(
+                f"Output {output_spec.arg.name} became a literal; changing output "
+                "representation is not supported"
+            )
+        if output_spec.arg.name != output.name:
+            output_spec.arg = replace(output_spec.arg, name=output.name)
+            modified = True
+    return modified
+
+
 class ExportedProgramPassBase(ABC):
     """
     Base interface for implementing passes that operate on ExportedProgram.
@@ -245,12 +290,19 @@ class ExportedProgramPassBase(ABC):
     def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
         """
         Runs the precondition check, the pass itself, and the postcondition check.
+
+        Output specs are realigned with the graph between the pass and the
+        postcondition check, so ``ensures`` sees a self-consistent program.
         """
 
         self.requires(exported_program)
         res = self.call(exported_program)
-        self.ensures(exported_program)
-        return res
+        signature_modified = _sync_output_specs(res.exported_program)
+        self.ensures(res.exported_program)
+        return ExportedProgramPassResult(
+            res.exported_program,
+            res.modified or signature_modified,
+        )
 
     @abstractmethod
     def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:

exir/tests/test_pass_infra.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 # Copyright 2026 Arm Limited and/or its affiliates.
@@ -579,3 +579,62 @@
         new_input = self._find_input_node(new_graph_module)
 
         self.assertNotEqual(self._symbolic_input_shape(new_input), original_snapshot)
+
+
+class ExportedProgramPassBaseOutputSpecTest(unittest.TestCase):
+    """__call__ realigns output specs with the graph before ensures() runs."""
+
+    class _Model(torch.nn.Module):
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x + x
+
+    def _program(self) -> ExportedProgram:
+        return to_edge(
+            export(self._Model(), (torch.randn(2, 2),))
+        ).exported_program()
+
+    def test_replacing_the_output_node_updates_the_signature(self) -> None:
+        class ReplaceOutputPass(ExportedProgramPassBase):
+            def call(self, ep: ExportedProgram) -> ExportedProgramPassResult:
+                graph = ep.graph_module.graph
+                output_node = graph.output_node()
+                (old,) = output_node.args[0]
+                with graph.inserting_before(output_node):
+                    new = graph.call_function(
+                        exir_ops.edge.aten.mul.Tensor, (old.args[0], old.args[1])
+                    )
+                new.meta = dict(old.meta)
+                output_node.args = ((new,),)
+                return ExportedProgramPassResult(ep, True)
+
+        program = self._program()
+        result = ReplaceOutputPass()(program)
+
+        (spec,) = result.exported_program.graph_signature.output_specs
+        graph_output_name = result.exported_program.graph.output_node().args[0][0].name
+        self.assertEqual(spec.arg.name, graph_output_name)
+        result.exported_program.validate()
+
+    def test_a_signature_only_change_is_reported_as_modified(self) -> None:
+        """A pass that renames the output reports modified even if it says False."""
+
+        class RenameOutputPass(ExportedProgramPassBase):
+            def call(self, ep: ExportedProgram) -> ExportedProgramPassResult:
+                ep.graph.output_node().args[0][0].name = "renamed_output"
+                return ExportedProgramPassResult(ep, False)
+
+        result = RenameOutputPass()(self._program())
+
+        self.assertTrue(result.modified)
+        (spec,) = result.exported_program.graph_signature.output_specs
+        self.assertEqual(spec.arg.name, "renamed_output")
+
+    def test_output_count_mismatch_is_rejected(self) -> None:
+        class DropOutputPass(ExportedProgramPassBase):
+            def call(self, ep: ExportedProgram) -> ExportedProgramPassResult:
+                output_node = ep.graph.output_node()
+                output_node.args = ((*output_node.args[0], output_node.args[0][0]),)
+                return ExportedProgramPassResult(ep, True)
+
+        with self.assertRaisesRegex(ExportPassBaseError, "output specs"):
+            DropOutputPass()(self._program())