From e61ec05b25e0c1bf31eb0bb87fd05d5c46e8be67 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 25 Apr 2025 16:48:26 -0700 Subject: [PATCH 1/2] Update type annotations for passes More clear on the documentation site. --- onnxscript/ir/passes/_pass_infra.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/onnxscript/ir/passes/_pass_infra.py b/onnxscript/ir/passes/_pass_infra.py index 56566e7556..165e2350e5 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 __all__ = [ "PassBase", @@ -180,11 +180,13 @@ class InPlacePass(PassBase): """A pass that modifies the input model in place and returns it.""" @property - def in_place(self) -> bool: + def in_place(self) -> Literal[True]: + """An in-place pass is in place.""" return True @property - def changes_input(self) -> bool: + def changes_input(self) -> Literal[True]: + """An in-place pass changes the input model.""" return True @@ -192,11 +194,13 @@ class FunctionalPass(PassBase): """A pass that returns a new model but does not modify the input model.""" @property - def in_place(self) -> bool: + def in_place(self) -> Literal[False]: + """A functional pass is not in place.""" return False @property - def changes_input(self) -> bool: + def changes_input(self) -> Literal[False]: + """A functional pass does not change the input model.""" return False From ea2f042010ade017e76bd74f970070d90fcb851f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 25 Apr 2025 20:18:38 -0700 Subject: [PATCH 2/2] final --- onnxscript/ir/passes/_pass_infra.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/onnxscript/ir/passes/_pass_infra.py b/onnxscript/ir/passes/_pass_infra.py index 165e2350e5..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 Literal, Sequence +from typing import Literal, Sequence, final __all__ = [ "PassBase", @@ -180,11 +180,13 @@ class InPlacePass(PassBase): """A pass that modifies the input model in place and returns it.""" @property + @final def in_place(self) -> Literal[True]: """An in-place pass is in place.""" return True @property + @final def changes_input(self) -> Literal[True]: """An in-place pass changes the input model.""" return True @@ -194,11 +196,13 @@ class FunctionalPass(PassBase): """A pass that returns a new model but does not modify the input model.""" @property + @final def in_place(self) -> Literal[False]: """A functional pass is not in place.""" return False @property + @final def changes_input(self) -> Literal[False]: """A functional pass does not change the input model.""" return False