broadcast the two input shapes for transposed matmul

	modified:   py/torch_tensorrt/fx/converters/acc_ops_converters.py
	modified:   py/torch_tensorrt/fx/test/passes/test_fuse_permute_matmul_trt.py

Author: XingFei Xi

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -37,6 +37,14 @@ def trt_transposed_matmul_converter(network, target, args, kwargs, name):
         lhs = get_trt_tensor(network, lhs, f"{name}_lhs")
     if isinstance(rhs, torch.nn.Parameter):
         rhs = get_trt_tensor(network, rhs, f"{name}_rhs")
+
+    lhs, rhs = broadcast(
+        network,
+        lhs,
+        rhs,
+        f"{lhs.name}_broadcast",
+        f"{rhs.name}_broadcast",
+    )
     layer = network.add_matrix_multiply(
         lhs,
         trt.MatrixOperation.TRANSPOSE if lhs_transposed else trt.MatrixOperation.NONE,

py/torch_tensorrt/fx/test/passes/test_fuse_permute_matmul_trt.py

@@ -22,6 +22,7 @@ def permute021(x):
 class TestFusePermuteMatmul(AccTestCase):
     @parameterized.expand(
         [
+            ("transpose_lhs_bmm_broadcast", (3, 2), (3, 3, 4), tranpose_last_two_dims),
             ("transpose_lhs_bmm", (3, 3, 2), (3, 3, 4), tranpose_last_two_dims),
             param(
                 "transpose_rhs_bmm", (3, 2, 3), (3, 4, 3), rhs_op=tranpose_last_two_dims
@@ -37,6 +38,8 @@ class TestFusePermuteMatmul(AccTestCase):
                 lambda x: x.permute(0, 1, 3, 2),
                 torch.matmul,
             ),
+            param("transpose_lhs_bmm_broadcast", (3, 2), (3, 3, 4), tranpose_last_two_dims, op=torch.matmul),
+            param("transpose_rhs_bmm_broadcast", (3, 3, 4), (3, 4), rhs_op=tranpose_last_two_dims, op=torch.matmul),
         ]
     )
     def test_fuse_permute_matmul(
@@ -53,11 +56,13 @@ def forward(self, x, y):
                 return op(lhs_op(x), rhs_op(y))
 
         inputs = [torch.randn(*lhs_shape), torch.randn(*rhs_shape)]
+        print("!!!!!!!!! come")
         self.run_test(
             TestModule(),
             inputs,
             {trt_transposed_matmul},
             apply_passes=[fuse_permute_matmul],
+            test_implicit_batch_dim=(len(lhs_shape) == len(rhs_shape)),
         )
 
     @parameterized.expand(