Correcting pre-commit hooks

Author: apbose

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -87,7 +87,13 @@ def set_layer_name(layer: TRTLayer, target: Target, name: str, is_acc=True) -> N
             the node represents.
         name (str): Consists of fx node.name with optional suffix.
     """
-    target_name = target if isinstance(target, str) else f"acc_ops.{target.__name__}" if is_acc else f"aten_ops.{target.__name__}"
+    target_name = (
+        target
+        if isinstance(target, str)
+        else f"acc_ops.{target.__name__}"
+        if is_acc
+        else f"aten_ops.{target.__name__}"
+    )
     layer.name = f"[{layer.type.name}]-[{target_name}]-[{name}]"
 
 

py/torch_tensorrt/fx/converters/operator.py

@@ -1118,17 +1118,17 @@ def add_expand(network, target, kwargs, name):
 
     ranks = len(input_val.shape)
     # TRT does not support different dimension size
-    #though this condition is not seen in the case of bmm 
+    # though this condition is not seen in the case of bmm
     # where input_t and shape dimensions are not equal
     assert len(shape) >= ranks
-    if(len(shape) != ranks):
-            shape_tuple = tuple([0] * len(shape))
-            shape_tensor = get_trt_tensor(network, input_t, f"{name}_shape")
-            input_val, shape_tensor = broadcast(network, input_val, shape_tensor, 
-                                  f"{name}_input_val",
-                                  f"{name}_shape_val")
-            ranks = len(shape)
-            
+    if len(shape) != ranks:
+        shape_tuple = tuple([0] * len(shape))
+        shape_tensor = get_trt_tensor(network, input_t, f"{name}_shape")
+        input_val, shape_tensor = broadcast(
+            network, input_val, shape_tensor, f"{name}_input_val", f"{name}_shape_val"
+        )
+        ranks = len(shape)
+
     shape = [input_val.shape[i] if shape[i] == -1 else shape[i] for i in range(ranks)]
 
     inshape = tuple(input_val.shape)

py/torch_tensorrt/fx/passes/lower_basic_pass_aten.py

@@ -417,11 +417,11 @@ def compose_bmm(
             input_n = node.all_input_nodes[0]
             other_n = node.all_input_nodes[1]
 
-             # If no input nodes are available, the bmm argument itself could be an input
+            # If no input nodes are available, the bmm argument itself could be an input
             # Alternatively, if the node has no users, it can be eliminated
             if len(input_n.all_input_nodes) == 0 or len(node.users) == 0:
                 return PassResult(module, modified)
-            
+
             output = next(iter(node.users))
             input_input_n = input_n.all_input_nodes[0]
             if (
@@ -434,7 +434,7 @@ def compose_bmm(
                     + "Skipping bmm lowering on this operation"
                 )
                 return PassResult(module, modified)
-            
+
             real_input = input_input_n.all_input_nodes[0]
             input_other_n = other_n.all_input_nodes[0]
             if (
@@ -447,7 +447,7 @@ def compose_bmm(
                     + "Skipping bmm lowering on this operation"
                 )
                 return PassResult(module, modified)
-            
+
             real_other = input_other_n.all_input_nodes[0]
             if len(real_other.meta["val"].size()) == 2:
                 new_func = aten_compose_bmm_2d

py/torch_tensorrt/fx/test/converters/aten_op/test_matmul_aten.py

@@ -19,13 +19,13 @@ class TestMatMulConverter(DispatchTestCase):
         [
             ("2_2", (2, 3), (3, 2)),
             ("2_2", (2, 3), (3, 1)),
-            #FIXME torch.ops.aten.mv.default for (2,3), (3,1) - should mv be lowered to mm?
+            # FIXME torch.ops.aten.mv.default for (2,3), (3,1) - should mv be lowered to mm?
             # (2,3), (3,) torch.ops.aten.mv.default
-            # Following cases use torch.ops.aten.bmm.defauly 
-            #("4_3", (3,1,3,2), (2,2,3)),
-            #("3_4", (3,1,3,2), (2,2,3)),
-            #("3_4", (2, 2, 3), (3, 1, 3, 3)),
-            #("4_2", (1, 2, 2, 3), (3, 2)),
+            # Following cases use torch.ops.aten.bmm.defauly
+            # ("4_3", (3,1,3,2), (2,2,3)),
+            # ("3_4", (3,1,3,2), (2,2,3)),
+            # ("3_4", (2, 2, 3), (3, 1, 3, 3)),
+            # ("4_2", (1, 2, 2, 3), (3, 2)),
         ]
     )
     def test_matmul_other_constant(self, _, input_shape, other_shape):
@@ -38,7 +38,7 @@ def forward(self, input):
                 return torch.matmul(input, self.other)
 
         inputs = [torch.randn(*input_shape)]
-        
+
         self.run_test(
             MatMul(),
             inputs,
@@ -50,14 +50,13 @@ def forward(self, input):
         [
             ("2_2", (2, 3), (3, 2)),
             ("1_2", (1, 3), (3, 2)),
-            #FIXME torch.ops.aten.mv.default for (2,3), (3,1) - should mv be lowered to mm?
+            # FIXME torch.ops.aten.mv.default for (2,3), (3,1) - should mv be lowered to mm?
             # (2,3), (3,) torch.ops.aten.mv.default
-            # Following cases use torch.ops.aten.bmm.defauly 
-            #("4_3", (3,1,3,2), (2,2,3)),
-            #("3_4", (3,1,3,2), (2,2,3)),
-            #("3_4", (2, 2, 3), (3, 1, 3, 3)),
-            #("4_2", (1, 2, 2, 3), (3, 2)),
-            
+            # Following cases use torch.ops.aten.bmm.defauly
+            # ("4_3", (3,1,3,2), (2,2,3)),
+            # ("3_4", (3,1,3,2), (2,2,3)),
+            # ("3_4", (2, 2, 3), (3, 1, 3, 3)),
+            # ("4_2", (1, 2, 2, 3), (3, 2)),
         ]
     )
     def test_matmul_input_constant(self, _, input_shape, other_shape):
@@ -75,8 +74,8 @@ def forward(self, other):
             MatMul(),
             inputs,
             expected_ops={torch.ops.aten.mm.default},
-            test_explicit_batch_dim=True 
-            #test_explicit_batch_dim=(len(other_shape) <= 2),
+            test_explicit_batch_dim=True
+            # test_explicit_batch_dim=(len(other_shape) <= 2),
         )
 
     @parameterized.expand(
@@ -96,7 +95,7 @@ def forward(self, input, other):
                 return torch.matmul(input, other)
 
         inputs = [torch.randn(*input_shape), torch.randn(*other_shape)]
-        test_explicit_batch_dim = not(
+        test_explicit_batch_dim = not (
             input_shape[0] == other_shape[0]
             and len(input_shape) > 2
             and len(other_shape) > 2
@@ -108,7 +107,8 @@ def forward(self, input, other):
             test_explicit_batch_dim=test_explicit_batch_dim,
         )
 
-    #FIXME: dynamic shape is giving bmm
+    # FIXME: dynamic shape is giving bmm
+
 
 if __name__ == "__main__":
-    run_tests()
+    run_tests()