Make pass_utils pyre-strict

Differential Revision: D70903186

Pull Request resolved: #9092

Author: mcremon-meta

backends/cadence/aot/pass_utils.py

@@ -7,7 +7,7 @@
 # pyre-strict
 
 from dataclasses import dataclass
-from typing import Callable, List, Optional, Set, Type, Union
+from typing import Callable, List, Optional, Set, Union
 
 import torch
 from executorch.backends.cadence.aot.utils import get_edge_overload_packet
@@ -32,33 +32,33 @@ class CadencePassAttribute:
 
 
 # A dictionary that maps an ExportPass to its attributes.
-ALL_CADENCE_PASSES: dict[Type[ExportPass], CadencePassAttribute] = {}
+ALL_CADENCE_PASSES: dict[ExportPass, CadencePassAttribute] = {}
 
 
-def get_cadence_pass_attribute(p: Type[ExportPass]) -> CadencePassAttribute:
+def get_cadence_pass_attribute(p: ExportPass) -> CadencePassAttribute:
     return ALL_CADENCE_PASSES[p]
 
 
 # A decorator that registers a pass.
 def register_cadence_pass(
     pass_attribute: CadencePassAttribute,
-) -> Callable[[Type[ExportPass]], Type[ExportPass]]:
-    def wrapper(cls: Type[ExportPass]) -> Type[ExportPass]:
+) -> Callable[[ExportPass], ExportPass]:
+    def wrapper(cls: ExportPass) -> ExportPass:
         ALL_CADENCE_PASSES[cls] = pass_attribute
         return cls
 
     return wrapper
 
 
-def get_all_available_cadence_passes() -> Set[Type[ExportPass]]:
+def get_all_available_cadence_passes() -> Set[ExportPass]:
     return set(ALL_CADENCE_PASSES.keys())
 
 
 # Create a new filter to filter out relevant passes from all passes.
 def create_cadence_pass_filter(
     opt_level: int, debug: bool = False
-) -> Callable[[Type[ExportPass]], bool]:
-    def _filter(p: Type[ExportPass]) -> bool:
+) -> Callable[[ExportPass], bool]:
+    def _filter(p: ExportPass) -> bool:
         pass_attribute = get_cadence_pass_attribute(p)
         return (
             pass_attribute.opt_level is not None

backends/cadence/aot/passes.py

@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from typing import Any, cast, List, Optional, Type
+from typing import Any, List, Optional
 
 import torch
 import torch.fx
@@ -71,7 +71,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 Argument = Any  # pyre-ignore
 
 
-def get_passes_in_default_order() -> List[Type[PassType]]:
+def get_passes_in_default_order() -> List[ExportPass]:
     passes = [
         InitializePipeline,
         RemoveRedundantOps.passes,
@@ -95,9 +95,8 @@ def get_cadence_passes(
     passes = get_passes_in_default_order()
     pass_filter = create_cadence_pass_filter(opt_level)
     filtered_passes = [
+        # pyre-ignore[20]: Expect argument graph_module
         filtered_pass()
-        # pyre-fixme[6]: In call `filter.__new__` ... got `List[Type[typing.Callable[[GraphModule], Optional[PassResult]]]]`.
         for filtered_pass in list(filter(pass_filter, passes))
     ]
-    # The type checker can't infer the proper type of the list comprehension.
-    return cast(List[Optional[PassResult]], filtered_passes)
+    return filtered_passes

backends/cadence/aot/replace_ops.py

@@ -1719,6 +1719,7 @@ def call_operator(self, op, args, kwargs, meta):
         )
 
 
+# pyre-ignore[6]: Incompatible parameter type (doesn't get the inheritance)
 register_cadence_pass(CadencePassAttribute(opt_level=0))(ReplaceScalarWithTensorArgPass)