diff --git a/onnxscript/ir/passes/_pass_infra.py b/onnxscript/ir/passes/_pass_infra.py index 56566e7556..18e5c8715b 100644 --- a/onnxscript/ir/passes/_pass_infra.py +++ b/onnxscript/ir/passes/_pass_infra.py @@ -16,7 +16,7 @@ import dataclasses import logging -from typing import Sequence +from typing import Literal, Sequence, final __all__ = [ "PassBase", @@ -180,11 +180,15 @@ class InPlacePass(PassBase): """A pass that modifies the input model in place and returns it.""" @property - def in_place(self) -> bool: + @final + def in_place(self) -> Literal[True]: + """An in-place pass is in place.""" return True @property - def changes_input(self) -> bool: + @final + def changes_input(self) -> Literal[True]: + """An in-place pass changes the input model.""" return True @@ -192,11 +196,15 @@ class FunctionalPass(PassBase): """A pass that returns a new model but does not modify the input model.""" @property - def in_place(self) -> bool: + @final + def in_place(self) -> Literal[False]: + """A functional pass is not in place.""" return False @property - def changes_input(self) -> bool: + @final + def changes_input(self) -> Literal[False]: + """A functional pass does not change the input model.""" return False