code: def test_permute_bmm_fusion(self):
      class TestModule(torch.nn.Module):
          def __init__(self, batch: int, k: int, n: int):
              super().__init__()
              self.other = torch.randn(batch, k, n)
  
          def forward(self, input: torch.Tensor):
              input1 = input.permute(0, 2, 1)
              output = torch.bmm(input1, self.other)
              return output
  
      batch, m, k, n = 6, 16, 8, 4
  
      trace_func = chain_passes(
          torch.fx.symbolic_trace, permute_matmul_fusion
      )
      module = TestModule(batch, k, n).eval()
      input = torch.randn(batch, k, m)
      traced = trace_func(module, [input])
      num_bmm = count_call_function(traced, torch.bmm)
      num_transpose_matmul = count_call_function(
          traced, transpose_matmul
      )
      self.assertEqual(num_bmm, 0)
      self.assertEqual(num_transpose_matmul, 1)
  
      self.assertTrue(torch.allclose(module(input), traced(input)))
  
file_path: pytorch/test/inductor/test_fx_fusion.py
imports: from typing import Any, Callable
  import torch
  from torch._inductor.fx_passes.pre_grad import linear_permute_fusion, linear_transpose, permute_linear_fusion, permute_matmul_fusion, sink_cat_after_pointwise, transpose_linear, transpose_matmul
  from torch.fx.passes.shape_prop import ShapeProp
  from torch.testing._internal.common_utils import run_tests, TestCase
label_desc: Use image node_id 4 for calling the TestFxFusion obj's underlying member method code with example usage: obj.test_permute_bmm_fusion() and returns: output