[Migration][DO NOT MERGE] Fix linting

ghstack-source-id: 8f46c5d
Pull Request resolved: #1333

Author: BowenBao

.lintrunner.toml

@@ -46,6 +46,21 @@ exclude_patterns = [
     'tests/**', # Skip linting test files for speed
     'onnxscript/**/*_test.py', # Skip linting test files for speed
     'onnxscript/function_libs/torch_lib/ops/**', # Operators typing do not play well with mypy
+    'onnxscript/optimizer/evaluator.py',  # FIXME
+    'onnxscript/optimizer/constant_folding.py',  # FIXME
+    'onnxscript/_legacy_ir/__init__.py',  # FIXME
+    'onnxscript/rewriter/onnxruntime/transformers/fastgelu.py',  # FIXME
+    'onnxscript/rewriter/onnxruntime/instance_to_group_normalization.py',  # FIXME
+    'onnxscript/rewriter/function_rule.py',  # FIXME
+    'onnxscript/_legacy_ir/irbuilder.py',  # FIXME
+    'onnxscript/optimizer/fold_constants_v0.py',  # FIXME
+    'onnxscript/rewriter/pattern.py',  # FIXME
+    'onnxscript/rewriter/onnxruntime/transformers/multihead_attention.py',  # FIXME
+    'onnxscript/tools/function_unittest_producer.py',  # FIXME
+    'onnxscript/_legacy_ir/visitor.py',  # FIXME
+    'onnxscript/_legacy_ir/protobuilder.py',  # FIXME
+    'onnxscript/rewriter/onnxruntime/transformers/layernorm.py',  # FIXME
+    'onnxscript/ir/serde.py',  # FIXME
 ]
 command = [
     'python',
@@ -108,6 +123,9 @@ exclude_patterns = [
     'tests/functions/**',
     'tests/models/**',
     'tests/onnx_backend_test_code/**',
+    'onnxscript/optimizer/**',  # FIXME
+    'onnxscript/rewriter/**',  # FIXME
+    'onnxscript/_legacy_ir/**',  # FIXME
 ]
 command = [
     'python',

onnxscript/_legacy_ir/__init__.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-
 import dataclasses
 from collections import deque
 from typing import List, Tuple, Union

onnxscript/_legacy_ir/irbuilder.py

@@ -36,9 +36,7 @@ def visit_model(self, model_proto: onnx.ModelProto) -> ir.Model:
         self._function_shape_env.load_from_model_proto(model_proto)
         self._ir_version = model_proto.ir_version
         version_map = {x.domain: x.version for x in model_proto.opset_import}
-        functions = [
-            self.visit_function(function) for function in model_proto.functions
-        ]
+        functions = [self.visit_function(function) for function in model_proto.functions]
         self.functions = {function.id: function for function in functions}
         graph = self.visit_graph(model_proto.graph)
         model = ir.Model()
@@ -148,13 +146,9 @@ def process_node(self, node):
             node_ir.attributes[attr.name] = attr_val
         # Set constant-value for Constant node:
         if node.op_type == "Constant" and node.domain in {"", "ai.onnx"}:
-            node_ir.outputs[0].value = utils.get_constant_node_value(
-                node, node.output[0]
-            )
+            node_ir.outputs[0].value = utils.get_constant_node_value(node, node.output[0])
 
-    def process_attribute(
-        self, attr: onnx.AttributeProto
-    ) -> ir.Graph | list[ir.Graph] | Any:
+    def process_attribute(self, attr: onnx.AttributeProto) -> ir.Graph | list[ir.Graph] | Any:
         if attr.HasField("g"):
             return self.visit_graph(attr.g)
         elif len(attr.graphs) > 0:
@@ -188,9 +182,7 @@ def process_function_input(self, input: str):
     def process_function_output(self, output: str):
         value = self.lookup(output)
         if value is None:
-            print(
-                f"WARNING: Function contains no definition for output '{output.name}'."
-            )
+            print(f"WARNING: Function contains no definition for output '{output.name}'.")
         else:
             value.is_output = True
 
@@ -201,7 +193,7 @@ def process_value_info(self, value_info: onnx.ValueInfoProto):
             existing_value.identity_merge_from(ir_value)
             ir_value = existing_value
 
-        if self._ir_version >= 10:  # noqa: PLR2004
+        if self._ir_version >= 10:
             # ONNX >= 1.16 where value_info can be defined in function
             self.bind(ir_value.name, ir_value)
         elif function_id is not None:

onnxscript/_legacy_ir/protobuilder.py

@@ -91,9 +91,7 @@ def visit_ir_function(
             function_proto.value_info.append(val)
         return function_proto
 
-    def process_ir_node(
-        self, ir_node: ir.Node, node_proto: onnx.NodeProto
-    ) -> onnx.NodeProto:
+    def process_ir_node(self, ir_node: ir.Node, node_proto: onnx.NodeProto) -> onnx.NodeProto:
         node_proto.op_type = ir_node.op_type
         node_proto.domain = ir_node.domain
         # Copy over node properties

onnxscript/_legacy_ir/protobuilder_test.py

@@ -201,9 +201,7 @@ def test_com_microsoft_opset_is_supported_in_protobuilder(self):
                 onnx.helper.make_tensor(
                     "weight", onnx.TensorProto.FLOAT16, [320, 1, 1], weight
                 ),
-                onnx.helper.make_tensor(
-                    "bias", onnx.TensorProto.FLOAT16, [320, 1, 1], bias
-                ),
+                onnx.helper.make_tensor("bias", onnx.TensorProto.FLOAT16, [320, 1, 1], bias),
             ]
         )
         ir = irbuilder.build_ir(model)

onnxscript/_legacy_ir/visitor.py

@@ -70,7 +70,7 @@ def process_value_info(
         self, value_info: onnx.ValueInfoProto
     ) -> tuple[ir.FunctionId | None, ir.Value]:
         name = value_info.name
-        if len(splits := name.split("/")) == 2:  # noqa: PLR2004
+        if len(splits := name.split("/")) == 2:
             # Experimental function value info format.
             # To be deprecated after ONNX 1.16, where value_info is introduced in FunctionProto.
             function_id, value_name = splits
@@ -79,7 +79,7 @@ def process_value_info(
             # 'overload' is introduced in ONNX 1.16, consider it as empty string prior to that.
             # The code is for future proof, in case overload is encoded in this format.
             overload = ""
-            if len(splits) == 3:  # noqa: PLR2004
+            if len(splits) == 3:
                 overload = splits[2]
             function_id = (domain, function_name, overload)
         else:
@@ -96,19 +96,14 @@ def save_to_value_info(
         function_id = f"{domain}::{function_name}"
 
         if value.type is not None:
-            return onnx.helper.make_value_info(
-                f"{function_id}/{value.name}", value.type
-            )
+            return onnx.helper.make_value_info(f"{function_id}/{value.name}", value.type)
         return None
 
     def lookup(self, function: onnx.FunctionProto, value_name: str) -> ir.Value | None:
         """Lookup ir value of 'value_name' inside 'function'."""
         function_id = ir.get_function_id(function)
         function_values = self._function_values.get(function_id)
-        if (
-            function_values is None
-            or (ir_value := function_values.get(value_name)) is None
-        ):
+        if function_values is None or (ir_value := function_values.get(value_name)) is None:
             logger.debug(
                 "Lookup Missed %s torch symbolic value info in function %s::%s.",
                 value_name,
@@ -124,9 +119,7 @@ def lookup(self, function: onnx.FunctionProto, value_name: str) -> ir.Value | No
         )
         return ir_value
 
-    def bind(
-        self, value: ir.Value, domain: str, function_name: str, overload: str
-    ) -> None:
+    def bind(self, value: ir.Value, domain: str, function_name: str, overload: str) -> None:
         """Bind ir value 'value' to 'value_name' inside 'function'."""
         function_id = (domain, function_name, overload)
         self._function_values.setdefault(function_id, {})[value.name] = value
@@ -309,9 +302,7 @@ def enter_function_scope(self, function: onnx.FunctionProto) -> None:
 
     def exit_function_scope(self) -> SubScope:
         sub_scope = self.current_scope().exit_sub_scope()
-        assert isinstance(
-            sub_scope.owner, onnx.FunctionProto
-        ), "Expected function scope."
+        assert isinstance(sub_scope.owner, onnx.FunctionProto), "Expected function scope."
         self._scopes.pop()
         return sub_scope
 
@@ -483,9 +474,7 @@ def input_element_type(self, node: onnx.NodeProto, index: int) -> int | None:
         info = self.get_input(node, index)
         return info.element_type if info is not None else None
 
-    def input_shape(
-        self, node: onnx.NodeProto, index: int
-    ) -> onnx.TensorShapeProto | None:
+    def input_shape(self, node: onnx.NodeProto, index: int) -> onnx.TensorShapeProto | None:
         info = self.get_input(node, index)
         return info.tensor_shape_proto() if info is not None else None
 
@@ -570,13 +559,11 @@ def process_node(self, node: onnx.NodeProto) -> Sequence[onnx.NodeProto] | None:
             # TODO: handle optional inputs
             def get_constant_value(i: int) -> onnx.TensorProto | None:
                 value = self.input_const_value(node, i)
-                if isinstance(value, np.ndarray) and value.size < 20:  # noqa: PLR2004
+                if isinstance(value, np.ndarray) and value.size < 20:
                     return onnx.numpy_helper.from_array(value, node.input[i])
                 return None
 
-            input_types = {
-                x: self.input_type(node, i) for i, x in enumerate(node.input)
-            }
+            input_types = {x: self.input_type(node, i) for i, x in enumerate(node.input)}
             input_data = {x: get_constant_value(i) for i, x in enumerate(node.input)}
             input_data = {k: v for k, v in input_data.items() if v is not None}
             if any(t is None for t in input_types.values()):
@@ -593,7 +580,7 @@ def get_constant_value(i: int) -> onnx.TensorProto | None:
                     output_types = onnx.shape_inference.infer_node_outputs(
                         schema, node, input_types, input_data
                     )
-                except Exception as e:  # noqa: BLE001
+                except Exception as e:
                     logger.debug(
                         "Skipping shape inference for node %s due to exception: %s",
                         node.name,
@@ -854,9 +841,7 @@ def process_function_node(
 
         self.enter_function_scope(mutable_function)
         if logger.level <= logging.INFO:
-            printable_actual_input_value_infos = [
-                str(x) for x in actual_input_value_infos
-            ]
+            printable_actual_input_value_infos = [str(x) for x in actual_input_value_infos]
             logger.info(
                 "Actual input value infos: %s",
                 printable_actual_input_value_infos,

onnxscript/function_libs/torch_lib/graph_building.py

@@ -828,7 +828,7 @@ def _override_with_symbolic_value_info_proto(self, onnx_model: onnx.ModelProto):
             new_value_info.pop(input.name, None)
         for output in onnx_model.graph.output:
             new_value_info.pop(output.name, None)
-        for tensor in onnx_model.graph.initializer:
+        for tensor in onnx_model.graph.initializer:  # type: ignore[assignment]
             new_value_info.pop(tensor.name, None)
         existing_value_info.update(new_value_info)
         onnx_model.graph.value_info.extend(existing_value_info.values())

onnxscript/ir/__init__.py

@@ -1,4 +1,5 @@
 """In-memory intermediate representation for ONNX graphs."""
+
 __all__ = [
     # Modules
     "serde",

onnxscript/ir/serde.py

@@ -55,7 +55,7 @@
 import logging
 import os
 import typing
-from typing import Any, Mapping, Sequence, List
+from typing import Any, List, Mapping, Sequence
 
 import numpy as np
 import onnx
@@ -431,7 +431,9 @@ def deserialize_function(proto: onnx.FunctionProto) -> _core.Function:
         doc_string=_get_field(proto, "doc_string"),
         opset_imports=deserialize_opset_import(proto.opset_import),
         name=(
-            f"{proto.name}_{proto.domain}" + f"__{proto.overload}" if hasattr(proto, "overload") and proto.overload else ""
+            f"{proto.name}_{proto.domain}" + f"__{proto.overload}"
+            if hasattr(proto, "overload") and proto.overload
+            else ""
         ),
     )
     attributes = [_deserialize_attribute(attr, []) for attr in proto.attribute_proto]

onnxscript/optimizer/constant_folding_test.py

@@ -117,13 +117,9 @@ def test_fold_inside_if_branch(self):
         )
         optimized = optimizer.optimize(model, num_iterations=1)
         self.assertEqual(len(optimized.graph.node), 1)
-        then_graph = onnx.helper.get_node_attr_value(
-            optimized.graph.node[0], "then_branch"
-        )
+        then_graph = onnx.helper.get_node_attr_value(optimized.graph.node[0], "then_branch")
         self.assertEqual(len(then_graph.node), 2)
-        else_graph = onnx.helper.get_node_attr_value(
-            optimized.graph.node[0], "else_branch"
-        )
+        else_graph = onnx.helper.get_node_attr_value(optimized.graph.node[0], "else_branch")
         self.assertEqual(len(else_graph.node), 2)
 
     def test_fold_if_propagate(self):
@@ -198,9 +194,7 @@ def test_fold_undefined_vars(self):
         """
         )
         # No optimizations expected. Just make sure it doesn't crash.
-        optimized = optimizer.optimize(
-            model, num_iterations=1, onnx_shape_inference=False
-        )
+        optimized = optimizer.optimize(model, num_iterations=1, onnx_shape_inference=False)
         self.assertEqual(len(optimized.graph.node), 6)
 
     def test_shape_inference(self):
@@ -336,9 +330,7 @@ def test_static_split_to_sequence_with_list_split_no_keepdims_and_squence_at_is_
         self.assertEqual(len(optimized.graph.node), 7)
         self.assertEqual(len(optimized.graph.node[1].output), 3)
         self.assertEqual(optimized.graph.node[1].op_type, "Split")
-        self.assertEqual(
-            len([n for n in optimized.graph.node if n.op_type == "Squeeze"]), 3
-        )
+        self.assertEqual(len([n for n in optimized.graph.node if n.op_type == "Squeeze"]), 3)
 
     def test_static_split_to_sequence_with_uneven_split(self):
         model = onnx.parser.parse_model(

onnxscript/optimizer/copy_propagation.py

@@ -56,11 +56,7 @@ def visit_node(self, node: onnx.NodeProto) -> None:
         if is_onnx_op(node, "ConcatFromSequence"):
             input = self.get_input(node, 0)
             new_axis = get_node_attr_value(node, "new_axis", 0)
-            if (
-                input is not None
-                and isinstance(input.symbolic_value, list)
-                and new_axis == 0
-            ):
+            if input is not None and isinstance(input.symbolic_value, list) and new_axis == 0:
                 node.op_type = "Concat"
                 node.input[:] = input.symbolic_value
                 for i in range(len(node.attribute)):
@@ -78,9 +74,7 @@ def do_copy_propagation(model: onnx.ModelProto, *, remove_unused: bool = True) -
         onnxscript.optimizer.remove_unused_nodes(model)
 
 
-def do_sequence_simplification(
-    model: onnx.ModelProto, *, remove_unused: bool = True
-) -> None:
+def do_sequence_simplification(model: onnx.ModelProto, *, remove_unused: bool = True) -> None:
     transformer = SymbolicEvaluator()
     transformer.visit_model(model)
     if remove_unused:

onnxscript/optimizer/evaluator.py

@@ -32,7 +32,7 @@ def get_evaluator(self, domain: str, op: str, version: int) -> callable | None:
         try:
             op_impl_class = onnx.reference.ops.load_op(domain, op, version)
             return op_impl_class.eval  # noqa: TRY300
-        except Exception:  # noqa: BLE001
+        except Exception:
             return None
 
     def evaluate(self, domain: str, op: str, version: int, *args, **kwargs) -> Any:
@@ -59,9 +59,7 @@ def get_input(self, node: onnx.NodeProto, index: int) -> ir.Value | None: ...
 
     def get_output(self, node: onnx.NodeProto, index: int) -> ir.Value | None: ...
 
-    def input_const_value(
-        self, node: onnx.NodeProto, index: int
-    ) -> ir.ConcreteValue: ...
+    def input_const_value(self, node: onnx.NodeProto, index: int) -> ir.ConcreteValue: ...
 
     def input_shape(
         self, node: onnx.NodeProto, index: int
@@ -75,9 +73,7 @@ def lookup_version(self, domain: str) -> int: ...
 
     def convert_attributes(self, attributes: Sequence[onnx.AttributeProto]) -> dict: ...
 
-    def new_constant(
-        self, name: str, value: Any
-    ) -> Sequence[onnx.NodeProto] | None: ...
+    def new_constant(self, name: str, value: Any) -> Sequence[onnx.NodeProto] | None: ...
 
 
 # A partial-evaluator function takes an IRContext and a node, and returns a list of
@@ -119,9 +115,7 @@ def __init__(self):
     def lookup_evaluators(self, domain: str, opname: str, version: int):
         evaluator_list = self.op_evaluators.get((domain, opname), [])
         return [
-            evaluator.function
-            for evaluator in evaluator_list
-            if evaluator.valid_for(version)
+            evaluator.function for evaluator in evaluator_list if evaluator.valid_for(version)
         ]
 
     def register(self, opname: str, domain: str = "", version=None):
@@ -427,9 +421,7 @@ def split_to_sequence(
 
 
 @register("SequenceAt")
-def sequence_at(
-    context: IRContext, node: onnx.NodeProto
-) -> Sequence[onnx.NodeProto] | None:
+def sequence_at(context: IRContext, node: onnx.NodeProto) -> Sequence[onnx.NodeProto] | None:
     input = context.get_input(node, 0)
     position = context.get_input(node, 1)
     output = context.get_output(node, 0)