Fix lint

Author: DenisVieriu97

backends/apple/mps/mps_preprocess.py

@@ -66,7 +66,7 @@ def preprocess(
             input_ids=[],
             output_ids=[],
             constant_ids=[],
-            graph_type=OpType.mps_graph
+            graph_type=OpType.mps_graph,
         )
 
         convert_model_to_fp16 = True
@@ -114,8 +114,13 @@ def handle_call_function(
     ) -> None:
         logging.info(f"Visiting: {node}, {node.target.__name__}")
 
-        if "delegation_tag" in node.meta and "metal_kernel" in node.meta["delegation_tag"]:
-            logging.info(f"Node '{node.target.__name__}' was marked as a Metal kernel by the MPSPartitioner!")
+        if (
+            "delegation_tag" in node.meta
+            and "metal_kernel" in node.meta["delegation_tag"]
+        ):
+            logging.info(
+                f"Node '{node.target.__name__}' was marked as a Metal kernel by the MPSPartitioner!"
+            )
             mps_graph.graph_type = OpType.metal_kernel
 
         if node.target.__name__ in node_visitors:

backends/apple/mps/operators/indexing_ops.py

@@ -13,13 +13,14 @@
 from executorch.backends.apple.mps.serialization.mps_graph_schema import (
     MPSEmbedding,
     MPSGraph,
-    MPSIndexTensor,
     MPSIndexPut,
     MPSIndexSelect,
+    MPSIndexTensor,
 )
 from executorch.backends.apple.mps.utils.mps_utils import get_input_node
-from executorch.exir.sym_util import eval_expr
 from executorch.backends.transforms import get_shape
+from executorch.exir.sym_util import eval_expr
+
 
 @register_node_visitor
 class IndexSelectVisitor(NodeVisitor):
@@ -41,6 +42,7 @@ def define_node(
 
         mps_graph.mps_nodes.append(mps_node)
 
+
 @register_node_visitor
 class IndexTensorVisitor(NodeVisitor):
     target = "aten.index.Tensor"
@@ -56,12 +58,13 @@ def define_node(
         mps_node = self.create_unary_node(node, mps_graph, MPSIndexTensor)
         tensors = cast(List[torch.fx.Node], node.args[1])
         for tensor in tensors:
-            mps_node.mpsnode_union.indices_id.append(self.define_tensor(tensor, mps_graph))
+            mps_node.mpsnode_union.indices_id.append(
+                self.define_tensor(tensor, mps_graph)
+            )
 
         mps_graph.mps_nodes.append(mps_node)
 
 
-
 # [MPS TODO]: Works on a single iteration of llama2, but subsequent tokens
 # are wrong when using Index put. Disabling it for now.
 @register_node_visitor
@@ -87,7 +90,6 @@ def infer_sizes(self, a: List[int], b: List[int]):
 
         return expandedSizes
 
-
     def define_node(
         self,
         node: torch.fx.Node,
@@ -103,13 +105,16 @@ def define_node(
 
         tensors = cast(List[torch.fx.Node], node.args[1])
         for tensor in tensors:
-            mps_node.mpsnode_union.indices_id.append(self.define_tensor(tensor, mps_graph))
+            mps_node.mpsnode_union.indices_id.append(
+                self.define_tensor(tensor, mps_graph)
+            )
 
         mps_node.mpsnode_union.values_id = self.define_tensor(
             get_input_node(node, 2), mps_graph
         )
         mps_graph.mps_nodes.append(mps_node)
 
+
 @register_node_visitor
 class EmbeddingVisitor(NodeVisitor):
     target = "aten.embedding.default"

backends/apple/mps/operators/unary_ops.py

@@ -30,6 +30,7 @@
     MPSLog,
     MPSLog10,
     MPSLog2,
+    MPSLogicalNot,
     MPSNeg,
     MPSReciprocal,
     MPSRound,
@@ -41,7 +42,6 @@
     MPSSqrt,
     MPSTan,
     MPSTanh,
-    MPSLogicalNot,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 

backends/apple/mps/partition/mps_partitioner.py

@@ -4,12 +4,13 @@
 #
 
 import logging
-from typing import cast, Any, Dict, List, Union
+from typing import Any, cast, Dict, List, Union
 
 import torch
 from executorch.backends.apple.mps.mps_preprocess import MPSBackend
 from executorch.backends.apple.mps.operators.node_visitor import get_node_visitors
 from executorch.backends.apple.mps.utils.mps_utils import is_parameter
+from executorch.backends.transforms import get_shape
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import (
     generate_partitions_from_list_of_nodes,
@@ -20,11 +21,10 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
-from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.backends.transforms import get_shape
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.DEBUG, format=FORMAT)
@@ -36,6 +36,7 @@
     exir_ops.edge.aten.index_put.default,
 ]
 
+
 class MPSOperatorSupport(OperatorSupportBase):
     def __init__(self, edge_program: torch.export.ExportedProgram, compiler_specs):
         self.node_visitors = get_node_visitors(edge_program)
@@ -90,7 +91,10 @@ def mps_graph_advanced_indexing_support(self, node: torch.fx.Node):
 
     def use_metal_kernel(self, node: torch.fx.Node):
         if node.target in METAL_KERNELS:
-            if node.target == exir_ops.edge.aten.index.Tensor or node.target == exir_ops.edge.aten.index_put.default:
+            if (
+                node.target == exir_ops.edge.aten.index.Tensor
+                or node.target == exir_ops.edge.aten.index_put.default
+            ):
                 if not self.mps_graph_advanced_indexing_support(node):
                     return True
         return False
@@ -104,7 +108,9 @@ def tag_nodes(self, partitions: List[Partition]) -> None:
                     logging.warning(f"[WARNING] Using Metal kernel for op {node.name}!")
                     # Partition the Metal kernel into a separate partition
                     crt_partition_counter += 1
-                    delegation_tag = f"{delegation_tag}_metal_kernel_{crt_partition_counter}"
+                    delegation_tag = (
+                        f"{delegation_tag}_metal_kernel_{crt_partition_counter}"
+                    )
                     crt_partition_counter += 1
                 else:
                     delegation_tag = f"{delegation_tag}_{crt_partition_counter}"

backends/apple/mps/serialization/mps_graph_schema.py

@@ -26,10 +26,12 @@ class MPSDataType(IntEnum):
     mps_data_type_complex_float16 = 10
     mps_data_type_complex_float32 = 11
 
+
 class OpType(IntEnum):
     mps_graph = 0
     metal_kernel = 1
 
+
 @dataclass
 class MPSNode1x1:
     input1_id: int
@@ -453,6 +455,7 @@ class MPSIndexPut(MPSNode1x1):
     values_shape: List[int] = field(default_factory=list)
     values_id: int = -1
 
+
 ##
 ## Shape ops
 ##