Fix lint

Author: DenisVieriu97

backends/apple/mps/test/test_mps.py

@@ -1190,7 +1190,11 @@ def forward(self, x, rows, columns):
         x = torch.arange(0, 12).resize(4, 3)
         rows = torch.tensor([[0, 0], [3, 3]])
         columns = torch.tensor([[0, 2], [0, 2]])
-        model_inputs = (x, rows, columns, )
+        model_inputs = (
+            x,
+            rows,
+            columns,
+        )
 
         self.lower_and_test_with_partitioner(
             module, model_inputs, func_name=inspect.stack()[0].function[5:]

backends/apple/mps/test/test_mps_indexing_ops.py

@@ -64,7 +64,9 @@ def forward(self, x):
                 return x[:, [0, 1, 0], [0, 1, 0]]
 
         module = IndexGet()
-        model_inputs = (torch.tensor([[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]),)
+        model_inputs = (
+            torch.tensor([[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [11, 12]]]),
+        )
 
         self.lower_and_test_with_partitioner(
             module, model_inputs, func_name=inspect.stack()[0].function[5:]
@@ -212,9 +214,12 @@ def forward(self, x, y, z):
         input = torch.ones(1, 8, 128, 8)
         indices = torch.tensor([1])
         values = torch.randn(8, 1, 8)
-        model_inputs = (input, indices, values, )
+        model_inputs = (
+            input,
+            indices,
+            values,
+        )
 
         self.lower_and_test_with_partitioner(
             module, model_inputs, func_name=inspect.stack()[0].function[5:]
         )
-

backends/apple/mps/test/test_mps_utils.py

@@ -22,13 +22,11 @@
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.exir.capture._config import ExecutorchBackendConfig
 from executorch.exir.tracer import Value
-from executorch.extension.pytree import tree_flatten
 from executorch.sdk import BundledProgram
 from executorch.sdk.bundled_program.config import MethodTestCase, MethodTestSuite
 from executorch.sdk.bundled_program.serialize import (
     serialize_from_bundled_program_to_flatbuffer,
 )
-from torch._export import capture_pre_autograd_graph
 from torch.export import export, ExportedProgram
 
 # Config for Capturing the weights, will be moved in the future
@@ -201,7 +199,7 @@ def lower_module_and_test_output(
         func_name: str,
         use_partitioner: bool = True,
         use_fp16: bool = False,
-        bundled_program = True,
+        bundled_program=True,
     ) -> ExirExportedProgram:
         """
         Helper testing function that takes a torch.nn.Module and lowers it to MPS with

examples/apple/mps/scripts/bench_utils.py

@@ -3,10 +3,12 @@
 #  Provided subject to the LICENSE file in the top level directory.
 #
 
-from torch._export.exported_program import ExportedProgram
-import torch
-import time
 import logging
+import time
+
+import torch
+from torch._export.exported_program import ExportedProgram
+
 
 def assert_outputs_equal(model_output, ref_output):
     """
@@ -19,16 +21,19 @@ def assert_outputs_equal(model_output, ref_output):
     # Compare the result from executor and eager mode direclty
     if isinstance(ref_output, tuple) or isinstance(ref_output, list):
         # Multiple outputs executor always returns tuple, even if there is one output
-        assert len(ref_output) == len(model_output), "Length of outputs is not matching!"
+        assert len(ref_output) == len(
+            model_output
+        ), "Length of outputs is not matching!"
         for i in range(len(ref_output)):
-            assert(
-                torch.allclose(
-                    model_output[i], ref_output[i], atol=1e-03, rtol=1e-03
-                )
+            assert torch.allclose(
+                model_output[i], ref_output[i], atol=1e-03, rtol=1e-03
             )
     else:
         # If one output, eager returns tensor while executor tuple of size 1
-        assert torch.allclose(model_output[0], ref_output, atol=1e-03, rtol=1e-03), "Outputs are not matching!"
+        assert torch.allclose(
+            model_output[0], ref_output, atol=1e-03, rtol=1e-03
+        ), "Outputs are not matching!"
+
 
 def bench_forward(func, *args):
     # warmup
@@ -41,39 +46,44 @@ def bench_forward(func, *args):
     end = time.time()
     return end - start
 
+
 def executorch_forward_pass(model, inputs):
     for _ in range(10):
         model.forward(inputs)
 
+
 def synchronize():
     torch.mps.synchronize()
 
+
 def pytorch_forward_pass(model, inputs):
     for _ in range(10):
         model(*inputs)
     synchronize()
 
+
 def get_mps_inputs(inputs):
     inputs_mps = []
     for tensor in inputs:
         inputs_mps.append(tensor.to("mps"))
     inputs_mps = tuple(inputs_mps)
     return inputs_mps
 
+
 def get_executorch_model(executorch_program: ExportedProgram):
     try:
         from executorch.extension.pybindings.portable_lib import (  # @manual
             _load_for_executorch_from_buffer,
         )
-        return  _load_for_executorch_from_buffer(
-            executorch_program.buffer
-        )
+
+        return _load_for_executorch_from_buffer(executorch_program.buffer)
     except ImportError:
         logging.info(
             "ExecuTorch MPS delegate was built without pybind support (not possible to run forward pass within python)"
         )
         return None
 
+
 def bench_torch(executorch_program: ExportedProgram, model, inputs, model_name):
     model = model.to("mps")
     inputs_mps = get_mps_inputs(inputs)
@@ -86,7 +96,10 @@ def bench_torch(executorch_program: ExportedProgram, model, inputs, model_name):
         logging.info(f"Model name: {model_name}")
         logging.info(f"Pytorch MPS forward pass: {t_pytorch} seconds")
         logging.info(f"ExecuTorch MPS forward pass: {t_executorch} seconds")
-        logging.info(f"ExecuTorch speedup: {((t_pytorch - t_executorch) / t_pytorch) * 100}%")
+        logging.info(
+            f"ExecuTorch speedup: {((t_pytorch - t_executorch) / t_pytorch) * 100}%"
+        )
+
 
 def compare_outputs(executorch_program: ExportedProgram, model, inputs, model_name):
     inputs_copy = []
@@ -99,4 +112,6 @@ def compare_outputs(executorch_program: ExportedProgram, model, inputs, model_na
     if executorch_model is not None:
         executorch_results = executorch_model.forward(inputs_copy)
         assert_outputs_equal(executorch_results, pytorch_results)
-        logging.info(F"Results between ExecuTorch forward pass with MPS backend and PyTorch forward pass are matching!")
+        logging.info(
+            f"Results between ExecuTorch forward pass with MPS backend and PyTorch forward pass for {model_name} are matching!"
+        )

examples/apple/mps/scripts/mps_example.py

@@ -10,6 +10,7 @@
 import logging
 
 import torch
+from examples.apple.mps.scripts.bench_utils import bench_torch, compare_outputs
 from executorch import exir
 from executorch.backends.apple.mps.mps_preprocess import MPSBackend
 from executorch.backends.apple.mps.partition.mps_partitioner import MPSPartitioner
@@ -27,10 +28,6 @@
 from executorch.sdk.bundled_program.serialize import (
     serialize_from_bundled_program_to_flatbuffer,
 )
-from examples.apple.mps.scripts.bench_utils import (
-    bench_torch,
-    compare_outputs,
-)
 
 from ....models import MODEL_NAME_TO_MODEL
 from ....models.model_factory import EagerModelFactory
@@ -40,7 +37,28 @@
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=FORMAT)
 
-if __name__ == "__main__":
+
+def get_bundled_program(executorch_program, example_inputs, expected_output):
+    method_test_suites = [
+        MethodTestSuite(
+            method_name="forward",
+            test_cases=[
+                MethodTestCase(
+                    inputs=example_inputs, expected_outputs=[expected_output]
+                )
+            ],
+        )
+    ]
+    logging.info(f"Expected output: {expected_output}")
+
+    bundled_program = BundledProgram(executorch_program, method_test_suites)
+    bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
+        bundled_program
+    )
+    return bundled_program_buffer
+
+
+def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "-m",
@@ -111,10 +129,10 @@
     )
 
     args = parser.parse_args()
+    return args
 
-    if args.model_name not in MODEL_NAME_TO_MODEL:
-        raise RuntimeError(f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}.")
 
+def get_model_config(args):
     model_config = {}
     model_config["module_name"] = MODEL_NAME_TO_MODEL[args.model_name][0]
     model_config["model_class_name"] = MODEL_NAME_TO_MODEL[args.model_name][1]
@@ -125,10 +143,17 @@
         if args.params:
             model_config["params"] = args.params
         model_config["use_kv_cache"] = True
+    return model_config
 
-    model, example_inputs, _ = EagerModelFactory.create_model(
-        **model_config
-    )
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    if args.model_name not in MODEL_NAME_TO_MODEL:
+        raise RuntimeError(f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}.")
+
+    model_config = get_model_config(args)
+    model, example_inputs, _ = EagerModelFactory.create_model(**model_config)
 
     model = model.eval()
     if args.check_correctness or args.bench_pytorch:
@@ -172,21 +197,9 @@
     model_name = f"{args.model_name}_mps"
 
     if args.bundled:
-        method_test_suites = [
-            MethodTestSuite(
-                method_name="forward",
-                test_cases=[
-                    MethodTestCase(
-                        inputs=example_inputs, expected_outputs=[model(*example_inputs)]
-                    )
-                ],
-            )
-        ]
-        logging.info(f"Expected output: {model(*example_inputs)}")
-
-        bundled_program = BundledProgram(executorch_program, method_test_suites)
-        bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
-            bundled_program
+        expected_output = model(*example_inputs)
+        bundled_program_buffer = get_bundled_program(
+            executorch_program, example_inputs, expected_output
         )
         model_name = f"{model_name}_bundled"
         extension = "fp16"