fix lints for decompose_cat

Author: mcr229

backends/xnnpack/_passes/__init__.py

@@ -17,11 +17,11 @@
 from executorch.backends.xnnpack._passes.convert_to_upsample_bilinear2d import (
     ConvertToUpsampleBilinear2d,
 )
+from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm_with_conv import (
     FuseBatchNormWithConvPass,
 )
-from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
 from executorch.backends.xnnpack._passes.remove_getitem_op import RemoveGetItemPass
 from executorch.backends.xnnpack._passes.tag_implicit_q_dq_pass import (

backends/xnnpack/_passes/decompose_cat.py

@@ -5,18 +5,17 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
+
 import torch
+from executorch.backends.xnnpack.utils.quant_utils import is_dequant, is_quant
 from executorch.exir.dialects._ops import ops as exir_ops
 
 from executorch.exir.pass_base import ExportPass, PassResult
-from executorch.backends.xnnpack.utils.quant_utils import (
-    is_dequant,
-    is_quant,
-)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARNING)
 
+
 class DecomposeConcatenate(ExportPass):
     """
     XNNPACK's Concatenate operation only supports concatenation for <= 5 tensors
@@ -25,37 +24,40 @@ class DecomposeConcatenate(ExportPass):
 
     Example:
     Before Pass:
-        cat: "f32" = torch.ops.aten.cat.default([t1, t2, t3, t4, t5, t6], 1); 
-    
+        cat: "f32" = torch.ops.aten.cat.default([t1, t2, t3, t4, t5, t6], 1);
+
     After Pass:
-        cat: "f32" = torch.ops.aten.cat.default([t1, t2, t3, t4, t5], 1); 
-        cat_1: "f32" = torch.ops.aten.cat.default([cat, t6], 1); 
+        cat: "f32" = torch.ops.aten.cat.default([t1, t2, t3, t4, t5], 1);
+        cat_1: "f32" = torch.ops.aten.cat.default([cat, t6], 1);
     """
 
     def call(self, graph_module: torch.fx.GraphModule):
         gm = graph_module
         for node in gm.graph.nodes:
-            if (node.op == "call_function" 
-                and node.target.__name__ == "aten.cat.default"):
+            if (
+                node.op == "call_function"
+                and node.target.__name__ == "aten.cat.default"
+            ):
                 concat_args = node.args
                 nodes_to_concat = node.args[0]
                 if len(nodes_to_concat) <= 5:
                     continue
-                
-                is_quantized = (all(is_dequant(node) for node in nodes_to_concat) 
-                                and all(is_quant(node) for node in node.users.keys()))
+
+                is_quantized = all(
+                    is_dequant(node) for node in nodes_to_concat
+                ) and all(is_quant(node) for node in node.users.keys())
 
                 # replace the cat args with the same args but only with the first 5 nodes
                 new_concat_args = (nodes_to_concat[:5],) + concat_args[1:]
-                node.args =  new_concat_args
+                node.args = new_concat_args
 
                 remainder_nodes_to_concat = nodes_to_concat[5:]
                 with gm.graph.inserting_after(node):
                     logger.debug(f"Decomposing cat node {node}")
                     remainder_concat_node = gm.graph.create_node(
                         "call_function",
                         target=exir_ops.edge.aten.cat.default,
-                        args=([],), # we will replace this remainder_nodes later
+                        args=([],),  # we will replace this remainder_nodes later
                         kwargs=node.kwargs,
                     )
                     node.replace_all_uses_with(remainder_concat_node)
@@ -64,11 +66,13 @@ def call(self, graph_module: torch.fx.GraphModule):
                     # concat node
                     q_params = nodes_to_concat[0].args[1:]
                     q_kwargs = nodes_to_concat[0].kwargs
-                    # Quantizer enforces all the inputs and output to a concat node must share 
+                    # Quantizer enforces all the inputs and output to a concat node must share
                     # the same qparams, this means the newly inserted q/dq pair must share the
                     # same qparams as the first quantized input in the concat node.
                     with gm.graph.inserting_after(node):
-                        logger.debug(f"Inserting Q/DQ pair for new cat node {remainder_concat_node}")
+                        logger.debug(
+                            f"Inserting Q/DQ pair for new cat node {remainder_concat_node}"
+                        )
                         q_node = gm.graph.create_node(
                             "call_function",
                             target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
@@ -82,10 +86,14 @@ def call(self, graph_module: torch.fx.GraphModule):
                             args=(q_node,) + q_params,
                             kwargs=q_kwargs,
                         )
-                    remainder_concat_node.args = ([dq_node] + remainder_nodes_to_concat,) + node.args[1:]
+                    remainder_concat_node.args = (
+                        [dq_node] + remainder_nodes_to_concat,
+                    ) + node.args[1:]
                 else:
-                    remainder_concat_node.args = ([node] + remainder_nodes_to_concat,) + node.args[1:]
-        
+                    remainder_concat_node.args = (
+                        [node] + remainder_nodes_to_concat,
+                    ) + node.args[1:]
+
         gm.recompile()
         new_gm = super().call(gm).graph_module
         return PassResult(new_gm, True)

backends/xnnpack/test/passes/test_decompose_cat_pass.py

@@ -4,17 +4,18 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import math
 import unittest
 
 import torch
-import math
 from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate
 from executorch.backends.xnnpack.test.tester import RunPasses, Tester
 
 
 class TestDecomposeCatPass(unittest.TestCase):
     PassStage = RunPasses([DecomposeConcatenate])
     cat_name = "executorch_exir_dialects_edge__ops_aten_cat_default"
+
     class Cat(torch.nn.Module):
         def forward(self, *args):
             xs = [*args]
@@ -29,9 +30,9 @@ def test_cat_gt_5(self):
             inputs = []
             for _ in range(num_inputs):
                 inputs.append(torch.randn(1, 2, 3))
-            
+
             num_cats = int(len(inputs) > 5)
-            num_cats += math.ceil((len(inputs) - 5)/4)
+            num_cats += math.ceil((len(inputs) - 5) / 4)
             (
                 Tester(self.Cat(), tuple(inputs))
                 .export()
@@ -42,7 +43,6 @@ def test_cat_gt_5(self):
                 .run_method_and_compare_outputs()
             )
 
-
     def test_cat_gt_10(self):
         inputs = [
             torch.randn(1, 2, 3),
@@ -51,9 +51,9 @@ def test_cat_gt_10(self):
             inputs = []
             for _ in range(num_inputs):
                 inputs.append(torch.randn(1, 2, 3))
-            
+
             num_cats = int(len(inputs) > 5)
-            num_cats += math.ceil((len(inputs) - 5)/4)
+            num_cats += math.ceil((len(inputs) - 5) / 4)
             (
                 Tester(self.Cat(), tuple(inputs))
                 .export()
@@ -72,9 +72,9 @@ def test_qs8_cat_gt_5(self):
             inputs = []
             for _ in range(num_inputs):
                 inputs.append(torch.randn(1, 2, 3))
-            
+
             num_cats = int(len(inputs) > 5)
-            num_cats += math.ceil((len(inputs) - 5)/4)
+            num_cats += math.ceil((len(inputs) - 5) / 4)
             (
                 Tester(self.Cat(), tuple(inputs))
                 .quantize()
@@ -86,7 +86,6 @@ def test_qs8_cat_gt_5(self):
                 .run_method_and_compare_outputs()
             )
 
-
     def test_cat_gt_10(self):
         inputs = [
             torch.randn(1, 2, 3),
@@ -95,9 +94,9 @@ def test_cat_gt_10(self):
             inputs = []
             for _ in range(num_inputs):
                 inputs.append(torch.randn(1, 2, 3))
-            
+
             num_cats = int(len(inputs) > 5)
-            num_cats += math.ceil((len(inputs) - 5)/4)
+            num_cats += math.ceil((len(inputs) - 5) / 4)
             (
                 Tester(self.Cat(), tuple(inputs))
                 .export()