From 9b9a6ed1ab08861921ebc3c2f9c81baf98079205 Mon Sep 17 00:00:00 2001
From: Mikael Sevenier <mikael.sevenier@sima.ai>
Date: Wed, 18 Mar 2020 19:30:05 -0700
Subject: [PATCH 1/6] fix conv transpose import from TF

---
 python/tvm/relay/frontend/tensorflow.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index bed32b7274af..24a1af19a40e 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -270,7 +270,7 @@ def _impl(inputs, attr, params):
                 attr['strides'][3], attr['strides'][1], attr['strides'][2]
             attr['data_format'] = 'NCHW'
 
-            if opname == 'conv_transpose' and len(attr['_output_shapes']) > 0:
+            if opname == 'conv_transpose' and len(attr['_output_shapes']) > 0 and attr['_output_shapes'][0] is not None:
                 tmp_shape = attr['_output_shapes'][0]
                 tmp_shape = [tmp_shape[ii] for ii in (0, 3, 1, 2)]
                 attr['_output_shapes'][0] = tmp_shape
@@ -355,7 +355,7 @@ def _impl(inputs, attr, params):
             kernel_h, kernel_w = attr['kernel_shape']
 
             pdata_shape = input_shape
-            if opname == 'conv_transpose' and len(attr['_output_shapes']) > 0:
+            if opname == 'conv_transpose' and len(attr['_output_shapes']) > 0 and attr['_output_shapes'][0] is not None:
                 pdata_shape = attr['_output_shapes'][0]
 
             if attr['data_format'] == 'NHWC':

From 1372b43fbfcc1b840091cdbcad80f398b8dbf1e4 Mon Sep 17 00:00:00 2001
From: Mikael Sevenier <mikael.sevenier@sima.ai>
Date: Tue, 3 Nov 2020 20:53:15 -0800
Subject: [PATCH 2/6] fix String::fromwe() to String::from()

---
 rust/tvm/src/lib.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rust/tvm/src/lib.rs b/rust/tvm/src/lib.rs
index 7e0682b86b33..27e794984094 100644
--- a/rust/tvm/src/lib.rs
+++ b/rust/tvm/src/lib.rs
@@ -53,7 +53,7 @@ macro_rules! export {
     ($($fn_name:expr),*) => {
         pub fn tvm_export(ns: &str) -> Result<(), tvm::Error> {
             $(
-                let name = String::fromwe(ns) + ::std::stringify!($fn_name);
+                let name = String::from(ns) + ::std::stringify!($fn_name);
                 tvm::runtime::function::register_override($fn_name, name, true)?;
             )*
             Ok(())

From aac6d884a39bdb3f3c68226f805627aebd38d881 Mon Sep 17 00:00:00 2001
From: Mikael Sevenier <mikael.sevenier@sima.ai>
Date: Tue, 23 May 2023 00:40:08 -0700
Subject: [PATCH 3/6] torch squeeze can use a list of axis

---
 python/tvm/relay/frontend/pytorch.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
index 2255396c0633..b1739241856c 100644
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -382,7 +382,10 @@ def squeeze(self, inputs, input_types):
             axis = None
         else:
             # TODO (t-vi): why is the cast to int needed? similarly elsewhere
-            axis = [int(inputs[1])]
+            # [MBS] squeeze axis can be a list
+            # axis = [int(inputs[1])]
+            inputs = [inputs[1]] if not isinstance(inputs[1], list) else inputs[1]
+            axis = [int(v) for v in inputs]
 
         return _op.transform.squeeze(data, axis)
 

From d98f224a0678f5b61a3437e02f19db68083a7acc Mon Sep 17 00:00:00 2001
From: Mikael Sevenier <mikael.sevenier@sima.ai>
Date: Wed, 24 May 2023 01:15:03 -0700
Subject: [PATCH 4/6] added test for squeeze with multiple axis (pytorch 2)

---
 tests/python/frontend/pytorch/test_forward.py | 150 ++++++------------
 1 file changed, 46 insertions(+), 104 deletions(-)

diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
index b602c14df3af..b84bfddbb2c2 100644
--- a/tests/python/frontend/pytorch/test_forward.py
+++ b/tests/python/frontend/pytorch/test_forward.py
@@ -198,9 +198,7 @@ def verify_model(
             if not tvm.runtime.enabled(target):
                 continue
             dev = tvm.device(target, 0)
-            exe = relay.create_executor(
-                kind, mod=mod, params=params, device=dev, target=target
-            ).evaluate()
+            exe = relay.create_executor(kind, mod=mod, params=params, device=dev, target=target).evaluate()
             result = exe(**compiled_input)
             if not isinstance(result, list):
                 result = [result]
@@ -578,9 +576,15 @@ class Squeeze2(Module):
         def forward(self, *args):
             return args[0].squeeze(1)
 
+    class Squeeze3(Module):
+        def forward(self, *args):
+            return args[0].squeeze((1, 3))
+
     input_data = torch.rand(input_shape).float()
     verify_model(Squeeze1().float().eval(), input_data=input_data)
     verify_model(Squeeze2().float().eval(), input_data=input_data)
+    if package_version.parse(torch.__version__) >= package_version.parse("2.0.0"):
+        verify_model(Squeeze3().float().eval(), input_data=input_data)
 
 
 @tvm.testing.uses_gpu
@@ -759,9 +763,7 @@ def test_forward_leakyrelu():
     verify_model(torch.nn.LeakyReLU().eval(), input_data=input_data)
     verify_model(torch.nn.LeakyReLU(negative_slope=0.05).eval(), input_data=input_data)
     verify_model(torch.nn.LeakyReLU(negative_slope=1.0, inplace=True).eval(), input_data=input_data)
-    verify_model(
-        torch.nn.LeakyReLU(negative_slope=1.25, inplace=True).eval(), input_data=input_data
-    )
+    verify_model(torch.nn.LeakyReLU(negative_slope=1.25, inplace=True).eval(), input_data=input_data)
 
 
 @tvm.testing.uses_gpu
@@ -959,13 +961,9 @@ def test_forward_maxpool3d():
         input_data = torch.rand(input_shape).float()
 
         verify_model(torch.nn.MaxPool3d(kernel_size=[1, 1, 1]).eval(), input_data)
-        verify_model(
-            torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[1, 2, 3]).eval(), input_data
-        )
+        verify_model(torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[1, 2, 3]).eval(), input_data)
         verify_model(torch.nn.MaxPool3d(kernel_size=[10, 10, 10]).eval(), input_data)
-        verify_model(
-            torch.nn.MaxPool3d(kernel_size=[4, 4, 4], padding=2, stride=2).eval(), input_data
-        )
+        verify_model(torch.nn.MaxPool3d(kernel_size=[4, 4, 4], padding=2, stride=2).eval(), input_data)
 
     # A functional variant (default strides = None case)
     class MaxPool3D(Module):
@@ -1034,9 +1032,7 @@ def forward(self, *args):
     input_data = torch.rand(input_shape).float()
     verify_model(torch.nn.AvgPool1d(kernel_size=[10]).eval(), input_data=input_data)
     verify_model(AvgPool1D2().float().eval(), input_data=input_data)
-    verify_model(
-        torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data
-    )
+    verify_model(torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data)
 
 
 @tvm.testing.uses_gpu
@@ -1052,16 +1048,12 @@ def forward(self, *args):
     input_data = torch.rand(input_shape).float()
     verify_model(torch.nn.AvgPool2d(kernel_size=[10, 10]).eval(), input_data=input_data)
     verify_model(AvgPool2D2().float().eval(), input_data=input_data)
-    verify_model(
-        torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data
-    )
+    verify_model(torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data)
 
     input_shape = [1, 1, 1, 9]
     input_data = torch.rand(input_shape).float()
     verify_model(
-        torch.nn.AvgPool2d(
-            kernel_size=[1, 2], stride=[1, 2], ceil_mode=True, count_include_pad=True
-        ).eval(),
+        torch.nn.AvgPool2d(kernel_size=[1, 2], stride=[1, 2], ceil_mode=True, count_include_pad=True).eval(),
         input_data=input_data,
     )
 
@@ -1079,9 +1071,7 @@ def forward(self, *args):
     input_data = torch.rand(input_shape).float()
     verify_model(torch.nn.AvgPool3d(kernel_size=[10, 10, 10]).eval(), input_data=input_data)
     verify_model(AvgPool3D1().float().eval(), input_data=input_data)
-    verify_model(
-        torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data
-    )
+    verify_model(torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data)
 
 
 @tvm.testing.uses_gpu
@@ -1178,9 +1168,7 @@ def forward(self, *args):
 @pytest.mark.parametrize("output_padding", [0, 1, 2], ids=lambda x: "output_padding=" + str(x))
 @pytest.mark.parametrize("groups", [1], ids=lambda x: "groups=" + str(x))
 @pytest.mark.parametrize("bias", [True, False], ids=lambda x: "bias=" + str(x))
-def test_forward_conv_transpose(
-    in_channels, out_channels, kernel_size, output_padding, bias, groups
-):
+def test_forward_conv_transpose(in_channels, out_channels, kernel_size, output_padding, bias, groups):
     """test_forward_conv_transpose"""
     # Note we do not test with groups  > 1 because that is not supported
     # in tvm for conv transpose operations
@@ -2331,9 +2319,7 @@ def test_forward_upsample3d():
     inp = torch.arange(1, 9, dtype=torch.float32).view(1, 1, 2, 2, 2)
     verify_model(torch.nn.Upsample(scale_factor=2, mode="nearest").eval(), inp)
     verify_model(torch.nn.Upsample(scale_factor=2, mode="trilinear").eval(), inp)
-    verify_model(
-        torch.nn.Upsample(scale_factor=2, mode="trilinear", align_corners=True).eval(), inp
-    )
+    verify_model(torch.nn.Upsample(scale_factor=2, mode="trilinear", align_corners=True).eval(), inp)
 
 
 def test_forward_nms():
@@ -2413,9 +2399,7 @@ def test_conv3d_transpose():
     for ishape in [(1, 8, 10, 5, 10), (1, 8, 5, 8, 8), (1, 8, 13, 7, 7)]:
         inp = torch.rand(ishape)
         verify_model(
-            torch.nn.ConvTranspose3d(
-                in_channels=8, out_channels=33, kernel_size=3, stride=2
-            ).eval(),
+            torch.nn.ConvTranspose3d(in_channels=8, out_channels=33, kernel_size=3, stride=2).eval(),
             inp,
         )
         verify_model(
@@ -2428,9 +2412,7 @@ def test_conv3d_transpose():
             ).eval(),
             inp,
         )
-        verify_model(
-            torch.nn.ConvTranspose3d(in_channels=8, out_channels=20, kernel_size=1).eval(), inp
-        )
+        verify_model(torch.nn.ConvTranspose3d(in_channels=8, out_channels=20, kernel_size=1).eval(), inp)
         verify_model(
             torch.nn.ConvTranspose3d(in_channels=8, out_channels=5, kernel_size=1, stride=2).eval(),
             inp,
@@ -2539,9 +2521,7 @@ def _impl(inputs, input_types):
             spatial_scale = inputs[2]
             pooled_size = (inputs[3], inputs[4])
             sampling_ratio = inputs[5]
-            return relay.op.vision.roi_align(
-                inputs[0], inputs[1], pooled_size, spatial_scale, sampling_ratio
-            )
+            return relay.op.vision.roi_align(inputs[0], inputs[1], pooled_size, spatial_scale, sampling_ratio)
 
         return _impl
 
@@ -2644,15 +2624,11 @@ def verify_model_vm(input_model, ishapes, idtype=None, idata=None, targets=None)
     # If no input_data provided, generate random data of specified dtype
     else:
         if idtype == torch.bool:
-            input_data = [
-                torch.Tensor.bool(torch.randint(low=0, high=2, size=shape)) for shape in ishapes
-            ]
+            input_data = [torch.Tensor.bool(torch.randint(low=0, high=2, size=shape)) for shape in ishapes]
         # Torch dtype can be float, complex, int, or Bool. Complex not supported,
         # so if not float or Bool, dtype must be int!
         elif not idtype.is_floating_point:
-            input_data = [
-                torch.randint(low=0, high=10, size=shape, dtype=idtype) for shape in ishapes
-            ]
+            input_data = [torch.randint(low=0, high=10, size=shape, dtype=idtype) for shape in ishapes]
         else:
             input_data = [torch.randn(shape, dtype=idtype) for shape in ishapes]
 
@@ -2826,6 +2802,7 @@ def forward(self, inp):
 @tvm.testing.uses_gpu
 def test_simple_rnn():
     """test_simple_rnn"""
+
     # The mixed tracing and scripting example from
     # https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html#mixing-scripting-and-tracing
     class DecisionGate(torch.nn.Module):
@@ -3949,12 +3926,8 @@ def forward(self, *args):
     # divisor could be either tensor or scalar
     divisor_tensor = torch.rand([5, 3]).float() + 0.5
     divisor_scalar = torch.tensor(1.0, dtype=torch.float32)
-    verify_model(
-        TrueDivide().float().eval(), input_data=[dividend, divisor_tensor], atol=1e-4, rtol=1e-4
-    )
-    verify_model(
-        TrueDivide().float().eval(), input_data=[dividend, divisor_scalar], atol=1e-4, rtol=1e-4
-    )
+    verify_model(TrueDivide().float().eval(), input_data=[dividend, divisor_tensor], atol=1e-4, rtol=1e-4)
+    verify_model(TrueDivide().float().eval(), input_data=[dividend, divisor_scalar], atol=1e-4, rtol=1e-4)
 
 
 @tvm.testing.uses_gpu
@@ -4030,6 +4003,7 @@ def test_weight_names():
 @tvm.testing.uses_gpu
 def test_duplicate_weight_use():
     """test_duplicate_weight_use"""
+
     # The test cases doesn't make any sense as a neural network,
     # the issue popped up in shared input/output embeddings of bert,
     # but this is quicker
@@ -4068,9 +4042,7 @@ def forward(self, *args):
     # vector x batched_matrix - 1D x ND
     tensor1 = torch.randn(5)
     tensor2 = torch.randn(2, 3, 5, 4)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # matrix x vector - 2D - 1D
     tensor1 = torch.randn(3, 4)
@@ -4085,51 +4057,37 @@ def forward(self, *args):
     # broadcasted matrix x batched matrix - 2D x ND
     tensor1 = torch.randn(10, 4)
     tensor2 = torch.randn(2, 3, 4, 5)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # batched matrix x vector - ND x 1D
     tensor1 = torch.randn(2, 3, 4, 5)
     tensor2 = torch.randn(5)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # batched matrix x broadcasted matrix - ND x 2D
     tensor1 = torch.randn(10, 3, 4)
     tensor2 = torch.randn(4, 5)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # batched matrix x batched matrix - ND x ND
     tensor1 = torch.randn(2, 10, 3, 4)
     tensor2 = torch.randn(2, 10, 4, 5)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # batched matrix x broadcasted matrix - ND x ND
     tensor1 = torch.randn(2, 5, 3, 4)
     tensor2 = torch.randn(2, 1, 4, 5)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # broadcasted matrix x batched matrix - ND x ND
     tensor1 = torch.randn(2, 1, 5, 4)
     tensor2 = torch.randn(2, 5, 4, 3)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
     # broadcasted matrix x broadcasted matrix - ND x ND
     tensor1 = torch.randn(3, 2, 3, 1, 5, 4)
     tensor2 = torch.randn(2, 1, 5, 4, 3)
-    verify_model(
-        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
-    )
+    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
 
 
 @pytest.mark.skip(reason="unsupported op aten::lift_fresh")
@@ -4284,6 +4242,7 @@ def forward(self, data):
 
 def test_forward_scatter():
     """test_forward_scatter"""
+
     # integer cannot be traced
     def test_fn_scatter(dim):
         return lambda data, index, src: torch.scatter(data, dim=dim, index=index, src=src)
@@ -4320,11 +4279,10 @@ def test_fn_scatter_add(dim):
 
 def test_forward_scatter_reduce():
     """test_forward_scatter_reduce"""
+
     # integer cannot be traced
     def test_fn_scatter_reduce(dim, reduce):
-        return lambda data, index, src: torch.scatter_reduce(
-            data, dim=dim, index=index, src=src, reduce=reduce
-        )
+        return lambda data, index, src: torch.scatter_reduce(data, dim=dim, index=index, src=src, reduce=reduce)
 
     in_data = torch.rand(3, 5) - 1
     in_index = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
@@ -4344,11 +4302,10 @@ def test_fn_scatter_reduce(dim, reduce):
 
 def test_forward_index_put():
     """test_forward_index_put"""
+
     # torch.index_put for 2D tensor and default accumulate (False)
     def test_fn_index_put2():
-        return lambda data, xidx, yidx, values: torch.index_put(
-            data, indices=[xidx, yidx], values=values
-        )
+        return lambda data, xidx, yidx, values: torch.index_put(data, indices=[xidx, yidx], values=values)
 
     # torch.index_put for 3D tensor and accumulate=True
     def test_fn_index_put3a():
@@ -4781,9 +4738,7 @@ def test_cross_entropy_loss():
     targets = torch.randint(0, 3, (N,))
     weights = torch.tensor([1, 2, 3]).float()
     verify_model(torch.nn.CrossEntropyLoss().eval(), input_data=[predictions, targets])
-    verify_model(
-        torch.nn.CrossEntropyLoss(weight=weights).eval(), input_data=[predictions, targets]
-    )
+    verify_model(torch.nn.CrossEntropyLoss(weight=weights).eval(), input_data=[predictions, targets])
 
     # class probabilities
     predictions = torch.randn(N, C).float()
@@ -4854,9 +4809,7 @@ def test_annotate_span():
     model = torchvision.models.resnet18().eval()
     inp = torch.randn([1, 3, 224, 224])
     trace = torch.jit.trace(model, inp).eval()
-    mod, _ = relay.frontend.from_pytorch(
-        trace, [("input", inp.shape)], use_parser_friendly_name=True
-    )
+    mod, _ = relay.frontend.from_pytorch(trace, [("input", inp.shape)], use_parser_friendly_name=True)
     relay.transform.AnnotateSpans()(mod)
 
 
@@ -5124,6 +5077,7 @@ def test_remainder(x, y):
 
 def test_softmax_fuse():
     """test_softmax_fuse"""
+
     # https://github.com/apache/tvm/issues/12001
     class Model(torch.nn.Module):
         """Pytorch model module"""
@@ -5157,11 +5111,7 @@ def forward(self, x):
         mod, params = relay.frontend.from_pytorch(torch.jit.trace(model, inp), [("inp0", sh)])
 
         with tvm.transform.PassContext(opt_level=4):
-            out = (
-                relay.create_executor("graph", mod, params=params)
-                .evaluate()(inp0=inp.numpy())
-                .numpy()
-            )
+            out = relay.create_executor("graph", mod, params=params).evaluate()(inp0=inp.numpy()).numpy()
 
         tvm.testing.assert_allclose(out, output_torch, rtol=1e-5, atol=1e-5)
 
@@ -5225,14 +5175,10 @@ def test_weight_norm():
     conv_wn = torch.nn.utils.weight_norm(torch.nn.Conv2d(in_channels, out_channels, kernel_size=3))
     verify_model(conv_wn.eval().float(), input_data_conv)
 
-    conv_wn_groups = torch.nn.utils.weight_norm(
-        torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, groups=2)
-    )
+    conv_wn_groups = torch.nn.utils.weight_norm(torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, groups=2))
     verify_model(conv_wn_groups.eval().float(), input_data_conv)
 
-    conv_wn = torch.nn.utils.weight_norm(
-        torch.nn.Conv2d(in_channels, out_channels, kernel_size=3), dim=1
-    )
+    conv_wn = torch.nn.utils.weight_norm(torch.nn.Conv2d(in_channels, out_channels, kernel_size=3), dim=1)
     verify_model(conv_wn.eval().float(), input_data_conv)
 
     linear_wn = torch.nn.utils.weight_norm(torch.nn.Linear(in_channels, out_channels))
@@ -5243,9 +5189,7 @@ def test_weight_norm():
 @tvm.testing.uses_gpu
 def test_baddbmm():
     def test_fn(alpha, beta):
-        return lambda inp, batch1, batch2: torch.baddbmm(
-            inp, batch1, batch2, beta=beta, alpha=alpha
-        )
+        return lambda inp, batch1, batch2: torch.baddbmm(inp, batch1, batch2, beta=beta, alpha=alpha)
 
     M = torch.randn(10, 3, 5)
     batch1 = torch.randn(10, 3, 4)
@@ -5270,9 +5214,7 @@ def forward(self, *args):
     shape_list = [(input_name, input_shape)]
     temp_dir = utils.tempdir().path
     script_module = torch.jit.trace(Conv2D(), [torch.rand(input_shape)])
-    _, _ = relay.frontend.from_pytorch(
-        script_module, shape_list, export_renamed_c_graph_path=temp_dir
-    )
+    _, _ = relay.frontend.from_pytorch(script_module, shape_list, export_renamed_c_graph_path=temp_dir)
 
     exported_c_graph_name = os.listdir(temp_dir)[0]
     assert "tvm_exported_c_graph_" in exported_c_graph_name

From 2ad163e835d75fe165777a6950a66dd254474b83 Mon Sep 17 00:00:00 2001
From: Mikael Sevenier <mikael.sevenier@sima.ai>
Date: Wed, 24 May 2023 01:17:22 -0700
Subject: [PATCH 5/6] clean old code

---
 python/tvm/relay/frontend/pytorch.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
index dc6640682c81..35cb86e41dfe 100644
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -399,8 +399,6 @@ def squeeze(self, inputs, input_types):
             axis = None
         else:
             # TODO (t-vi): why is the cast to int needed? similarly elsewhere
-            # [MBS] squeeze axis can be a list
-            # axis = [int(inputs[1])]
             inputs = [inputs[1]] if not isinstance(inputs[1], list) else inputs[1]
             axis = [int(v) for v in inputs]
 

From 9166241410543f869ee8ce8647bbc2a5760e8d2b Mon Sep 17 00:00:00 2001
From: Mikael Sevenier <mikael.sevenier@sima.ai>
Date: Wed, 24 May 2023 01:20:31 -0700
Subject: [PATCH 6/6] code without reformating

---
 tests/python/frontend/pytorch/test_forward.py | 144 +++++++++++++-----
 1 file changed, 104 insertions(+), 40 deletions(-)

diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
index b84bfddbb2c2..65134ca15a34 100644
--- a/tests/python/frontend/pytorch/test_forward.py
+++ b/tests/python/frontend/pytorch/test_forward.py
@@ -198,7 +198,9 @@ def verify_model(
             if not tvm.runtime.enabled(target):
                 continue
             dev = tvm.device(target, 0)
-            exe = relay.create_executor(kind, mod=mod, params=params, device=dev, target=target).evaluate()
+            exe = relay.create_executor(
+                kind, mod=mod, params=params, device=dev, target=target
+            ).evaluate()
             result = exe(**compiled_input)
             if not isinstance(result, list):
                 result = [result]
@@ -763,7 +765,9 @@ def test_forward_leakyrelu():
     verify_model(torch.nn.LeakyReLU().eval(), input_data=input_data)
     verify_model(torch.nn.LeakyReLU(negative_slope=0.05).eval(), input_data=input_data)
     verify_model(torch.nn.LeakyReLU(negative_slope=1.0, inplace=True).eval(), input_data=input_data)
-    verify_model(torch.nn.LeakyReLU(negative_slope=1.25, inplace=True).eval(), input_data=input_data)
+    verify_model(
+        torch.nn.LeakyReLU(negative_slope=1.25, inplace=True).eval(), input_data=input_data
+    )
 
 
 @tvm.testing.uses_gpu
@@ -961,9 +965,13 @@ def test_forward_maxpool3d():
         input_data = torch.rand(input_shape).float()
 
         verify_model(torch.nn.MaxPool3d(kernel_size=[1, 1, 1]).eval(), input_data)
-        verify_model(torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[1, 2, 3]).eval(), input_data)
+        verify_model(
+            torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[1, 2, 3]).eval(), input_data
+        )
         verify_model(torch.nn.MaxPool3d(kernel_size=[10, 10, 10]).eval(), input_data)
-        verify_model(torch.nn.MaxPool3d(kernel_size=[4, 4, 4], padding=2, stride=2).eval(), input_data)
+        verify_model(
+            torch.nn.MaxPool3d(kernel_size=[4, 4, 4], padding=2, stride=2).eval(), input_data
+        )
 
     # A functional variant (default strides = None case)
     class MaxPool3D(Module):
@@ -1032,7 +1040,9 @@ def forward(self, *args):
     input_data = torch.rand(input_shape).float()
     verify_model(torch.nn.AvgPool1d(kernel_size=[10]).eval(), input_data=input_data)
     verify_model(AvgPool1D2().float().eval(), input_data=input_data)
-    verify_model(torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data)
+    verify_model(
+        torch.nn.AvgPool1d(kernel_size=[5], stride=2, padding=2).eval(), input_data=input_data
+    )
 
 
 @tvm.testing.uses_gpu
@@ -1048,12 +1058,16 @@ def forward(self, *args):
     input_data = torch.rand(input_shape).float()
     verify_model(torch.nn.AvgPool2d(kernel_size=[10, 10]).eval(), input_data=input_data)
     verify_model(AvgPool2D2().float().eval(), input_data=input_data)
-    verify_model(torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data)
+    verify_model(
+        torch.nn.AvgPool2d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data
+    )
 
     input_shape = [1, 1, 1, 9]
     input_data = torch.rand(input_shape).float()
     verify_model(
-        torch.nn.AvgPool2d(kernel_size=[1, 2], stride=[1, 2], ceil_mode=True, count_include_pad=True).eval(),
+        torch.nn.AvgPool2d(
+            kernel_size=[1, 2], stride=[1, 2], ceil_mode=True, count_include_pad=True
+        ).eval(),
         input_data=input_data,
     )
 
@@ -1071,7 +1085,9 @@ def forward(self, *args):
     input_data = torch.rand(input_shape).float()
     verify_model(torch.nn.AvgPool3d(kernel_size=[10, 10, 10]).eval(), input_data=input_data)
     verify_model(AvgPool3D1().float().eval(), input_data=input_data)
-    verify_model(torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data)
+    verify_model(
+        torch.nn.AvgPool3d(kernel_size=5, stride=2, padding=2).eval(), input_data=input_data
+    )
 
 
 @tvm.testing.uses_gpu
@@ -1168,7 +1184,9 @@ def forward(self, *args):
 @pytest.mark.parametrize("output_padding", [0, 1, 2], ids=lambda x: "output_padding=" + str(x))
 @pytest.mark.parametrize("groups", [1], ids=lambda x: "groups=" + str(x))
 @pytest.mark.parametrize("bias", [True, False], ids=lambda x: "bias=" + str(x))
-def test_forward_conv_transpose(in_channels, out_channels, kernel_size, output_padding, bias, groups):
+def test_forward_conv_transpose(
+    in_channels, out_channels, kernel_size, output_padding, bias, groups
+):
     """test_forward_conv_transpose"""
     # Note we do not test with groups  > 1 because that is not supported
     # in tvm for conv transpose operations
@@ -2319,7 +2337,9 @@ def test_forward_upsample3d():
     inp = torch.arange(1, 9, dtype=torch.float32).view(1, 1, 2, 2, 2)
     verify_model(torch.nn.Upsample(scale_factor=2, mode="nearest").eval(), inp)
     verify_model(torch.nn.Upsample(scale_factor=2, mode="trilinear").eval(), inp)
-    verify_model(torch.nn.Upsample(scale_factor=2, mode="trilinear", align_corners=True).eval(), inp)
+    verify_model(
+        torch.nn.Upsample(scale_factor=2, mode="trilinear", align_corners=True).eval(), inp
+    )
 
 
 def test_forward_nms():
@@ -2399,7 +2419,9 @@ def test_conv3d_transpose():
     for ishape in [(1, 8, 10, 5, 10), (1, 8, 5, 8, 8), (1, 8, 13, 7, 7)]:
         inp = torch.rand(ishape)
         verify_model(
-            torch.nn.ConvTranspose3d(in_channels=8, out_channels=33, kernel_size=3, stride=2).eval(),
+            torch.nn.ConvTranspose3d(
+                in_channels=8, out_channels=33, kernel_size=3, stride=2
+            ).eval(),
             inp,
         )
         verify_model(
@@ -2412,7 +2434,9 @@ def test_conv3d_transpose():
             ).eval(),
             inp,
         )
-        verify_model(torch.nn.ConvTranspose3d(in_channels=8, out_channels=20, kernel_size=1).eval(), inp)
+        verify_model(
+            torch.nn.ConvTranspose3d(in_channels=8, out_channels=20, kernel_size=1).eval(), inp
+        )
         verify_model(
             torch.nn.ConvTranspose3d(in_channels=8, out_channels=5, kernel_size=1, stride=2).eval(),
             inp,
@@ -2521,7 +2545,9 @@ def _impl(inputs, input_types):
             spatial_scale = inputs[2]
             pooled_size = (inputs[3], inputs[4])
             sampling_ratio = inputs[5]
-            return relay.op.vision.roi_align(inputs[0], inputs[1], pooled_size, spatial_scale, sampling_ratio)
+            return relay.op.vision.roi_align(
+                inputs[0], inputs[1], pooled_size, spatial_scale, sampling_ratio
+            )
 
         return _impl
 
@@ -2624,11 +2650,15 @@ def verify_model_vm(input_model, ishapes, idtype=None, idata=None, targets=None)
     # If no input_data provided, generate random data of specified dtype
     else:
         if idtype == torch.bool:
-            input_data = [torch.Tensor.bool(torch.randint(low=0, high=2, size=shape)) for shape in ishapes]
+            input_data = [
+                torch.Tensor.bool(torch.randint(low=0, high=2, size=shape)) for shape in ishapes
+            ]
         # Torch dtype can be float, complex, int, or Bool. Complex not supported,
         # so if not float or Bool, dtype must be int!
         elif not idtype.is_floating_point:
-            input_data = [torch.randint(low=0, high=10, size=shape, dtype=idtype) for shape in ishapes]
+            input_data = [
+                torch.randint(low=0, high=10, size=shape, dtype=idtype) for shape in ishapes
+            ]
         else:
             input_data = [torch.randn(shape, dtype=idtype) for shape in ishapes]
 
@@ -2802,7 +2832,6 @@ def forward(self, inp):
 @tvm.testing.uses_gpu
 def test_simple_rnn():
     """test_simple_rnn"""
-
     # The mixed tracing and scripting example from
     # https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html#mixing-scripting-and-tracing
     class DecisionGate(torch.nn.Module):
@@ -3926,8 +3955,12 @@ def forward(self, *args):
     # divisor could be either tensor or scalar
     divisor_tensor = torch.rand([5, 3]).float() + 0.5
     divisor_scalar = torch.tensor(1.0, dtype=torch.float32)
-    verify_model(TrueDivide().float().eval(), input_data=[dividend, divisor_tensor], atol=1e-4, rtol=1e-4)
-    verify_model(TrueDivide().float().eval(), input_data=[dividend, divisor_scalar], atol=1e-4, rtol=1e-4)
+    verify_model(
+        TrueDivide().float().eval(), input_data=[dividend, divisor_tensor], atol=1e-4, rtol=1e-4
+    )
+    verify_model(
+        TrueDivide().float().eval(), input_data=[dividend, divisor_scalar], atol=1e-4, rtol=1e-4
+    )
 
 
 @tvm.testing.uses_gpu
@@ -4003,7 +4036,6 @@ def test_weight_names():
 @tvm.testing.uses_gpu
 def test_duplicate_weight_use():
     """test_duplicate_weight_use"""
-
     # The test cases doesn't make any sense as a neural network,
     # the issue popped up in shared input/output embeddings of bert,
     # but this is quicker
@@ -4042,7 +4074,9 @@ def forward(self, *args):
     # vector x batched_matrix - 1D x ND
     tensor1 = torch.randn(5)
     tensor2 = torch.randn(2, 3, 5, 4)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # matrix x vector - 2D - 1D
     tensor1 = torch.randn(3, 4)
@@ -4057,37 +4091,51 @@ def forward(self, *args):
     # broadcasted matrix x batched matrix - 2D x ND
     tensor1 = torch.randn(10, 4)
     tensor2 = torch.randn(2, 3, 4, 5)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # batched matrix x vector - ND x 1D
     tensor1 = torch.randn(2, 3, 4, 5)
     tensor2 = torch.randn(5)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # batched matrix x broadcasted matrix - ND x 2D
     tensor1 = torch.randn(10, 3, 4)
     tensor2 = torch.randn(4, 5)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # batched matrix x batched matrix - ND x ND
     tensor1 = torch.randn(2, 10, 3, 4)
     tensor2 = torch.randn(2, 10, 4, 5)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # batched matrix x broadcasted matrix - ND x ND
     tensor1 = torch.randn(2, 5, 3, 4)
     tensor2 = torch.randn(2, 1, 4, 5)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # broadcasted matrix x batched matrix - ND x ND
     tensor1 = torch.randn(2, 1, 5, 4)
     tensor2 = torch.randn(2, 5, 4, 3)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
     # broadcasted matrix x broadcasted matrix - ND x ND
     tensor1 = torch.randn(3, 2, 3, 1, 5, 4)
     tensor2 = torch.randn(2, 1, 5, 4, 3)
-    verify_model(MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"])
+    verify_model(
+        MatMul1().float().eval(), input_data=[tensor1, tensor2], expected_ops=["nn.batch_matmul"]
+    )
 
 
 @pytest.mark.skip(reason="unsupported op aten::lift_fresh")
@@ -4242,7 +4290,6 @@ def forward(self, data):
 
 def test_forward_scatter():
     """test_forward_scatter"""
-
     # integer cannot be traced
     def test_fn_scatter(dim):
         return lambda data, index, src: torch.scatter(data, dim=dim, index=index, src=src)
@@ -4279,10 +4326,11 @@ def test_fn_scatter_add(dim):
 
 def test_forward_scatter_reduce():
     """test_forward_scatter_reduce"""
-
     # integer cannot be traced
     def test_fn_scatter_reduce(dim, reduce):
-        return lambda data, index, src: torch.scatter_reduce(data, dim=dim, index=index, src=src, reduce=reduce)
+        return lambda data, index, src: torch.scatter_reduce(
+            data, dim=dim, index=index, src=src, reduce=reduce
+        )
 
     in_data = torch.rand(3, 5) - 1
     in_index = torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]])
@@ -4302,10 +4350,11 @@ def test_fn_scatter_reduce(dim, reduce):
 
 def test_forward_index_put():
     """test_forward_index_put"""
-
     # torch.index_put for 2D tensor and default accumulate (False)
     def test_fn_index_put2():
-        return lambda data, xidx, yidx, values: torch.index_put(data, indices=[xidx, yidx], values=values)
+        return lambda data, xidx, yidx, values: torch.index_put(
+            data, indices=[xidx, yidx], values=values
+        )
 
     # torch.index_put for 3D tensor and accumulate=True
     def test_fn_index_put3a():
@@ -4738,7 +4787,9 @@ def test_cross_entropy_loss():
     targets = torch.randint(0, 3, (N,))
     weights = torch.tensor([1, 2, 3]).float()
     verify_model(torch.nn.CrossEntropyLoss().eval(), input_data=[predictions, targets])
-    verify_model(torch.nn.CrossEntropyLoss(weight=weights).eval(), input_data=[predictions, targets])
+    verify_model(
+        torch.nn.CrossEntropyLoss(weight=weights).eval(), input_data=[predictions, targets]
+    )
 
     # class probabilities
     predictions = torch.randn(N, C).float()
@@ -4809,7 +4860,9 @@ def test_annotate_span():
     model = torchvision.models.resnet18().eval()
     inp = torch.randn([1, 3, 224, 224])
     trace = torch.jit.trace(model, inp).eval()
-    mod, _ = relay.frontend.from_pytorch(trace, [("input", inp.shape)], use_parser_friendly_name=True)
+    mod, _ = relay.frontend.from_pytorch(
+        trace, [("input", inp.shape)], use_parser_friendly_name=True
+    )
     relay.transform.AnnotateSpans()(mod)
 
 
@@ -5077,7 +5130,6 @@ def test_remainder(x, y):
 
 def test_softmax_fuse():
     """test_softmax_fuse"""
-
     # https://github.com/apache/tvm/issues/12001
     class Model(torch.nn.Module):
         """Pytorch model module"""
@@ -5111,7 +5163,11 @@ def forward(self, x):
         mod, params = relay.frontend.from_pytorch(torch.jit.trace(model, inp), [("inp0", sh)])
 
         with tvm.transform.PassContext(opt_level=4):
-            out = relay.create_executor("graph", mod, params=params).evaluate()(inp0=inp.numpy()).numpy()
+            out = (
+                relay.create_executor("graph", mod, params=params)
+                .evaluate()(inp0=inp.numpy())
+                .numpy()
+            )
 
         tvm.testing.assert_allclose(out, output_torch, rtol=1e-5, atol=1e-5)
 
@@ -5175,10 +5231,14 @@ def test_weight_norm():
     conv_wn = torch.nn.utils.weight_norm(torch.nn.Conv2d(in_channels, out_channels, kernel_size=3))
     verify_model(conv_wn.eval().float(), input_data_conv)
 
-    conv_wn_groups = torch.nn.utils.weight_norm(torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, groups=2))
+    conv_wn_groups = torch.nn.utils.weight_norm(
+        torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, groups=2)
+    )
     verify_model(conv_wn_groups.eval().float(), input_data_conv)
 
-    conv_wn = torch.nn.utils.weight_norm(torch.nn.Conv2d(in_channels, out_channels, kernel_size=3), dim=1)
+    conv_wn = torch.nn.utils.weight_norm(
+        torch.nn.Conv2d(in_channels, out_channels, kernel_size=3), dim=1
+    )
     verify_model(conv_wn.eval().float(), input_data_conv)
 
     linear_wn = torch.nn.utils.weight_norm(torch.nn.Linear(in_channels, out_channels))
@@ -5189,7 +5249,9 @@ def test_weight_norm():
 @tvm.testing.uses_gpu
 def test_baddbmm():
     def test_fn(alpha, beta):
-        return lambda inp, batch1, batch2: torch.baddbmm(inp, batch1, batch2, beta=beta, alpha=alpha)
+        return lambda inp, batch1, batch2: torch.baddbmm(
+            inp, batch1, batch2, beta=beta, alpha=alpha
+        )
 
     M = torch.randn(10, 3, 5)
     batch1 = torch.randn(10, 3, 4)
@@ -5214,7 +5276,9 @@ def forward(self, *args):
     shape_list = [(input_name, input_shape)]
     temp_dir = utils.tempdir().path
     script_module = torch.jit.trace(Conv2D(), [torch.rand(input_shape)])
-    _, _ = relay.frontend.from_pytorch(script_module, shape_list, export_renamed_c_graph_path=temp_dir)
+    _, _ = relay.frontend.from_pytorch(
+        script_module, shape_list, export_renamed_c_graph_path=temp_dir
+    )
 
     exported_c_graph_name = os.listdir(temp_dir)[0]
     assert "tvm_exported_c_graph_" in exported_c_graph_name