code: def permute_matmul_fusion(
      module: torch.fx.GraphModule,
  ) -> torch.fx.GraphModule:
      for node in module.graph.nodes:
          if node.op == "call_function" and (
              node.target == torch.bmm or node.target == torch.matmul
          ):
              normalized = NormalizedMatmulNode(node)
              input_A_node = normalized.get_input()
              input_B_node = normalized.get_other()
              input_A = input_A_node
              input_B = input_B_node
              Atrans = Btrans = False
              if (
                  input_A_node.op == "call_method"
                  and input_A_node.target == "permute"
                  and check_permute(input_A_node)
              ):
                  Atrans = True
                  if len(input_A_node.args) > 0:
                      input_A = input_A_node.args[0]
                  else:
                      input_A = input_A_node.kwargs["input"]
  
              if (
                  input_B_node.op == "call_method"
                  and input_B_node.target == "permute"
                  and check_permute(input_B_node)
              ):
                  Btrans = True
                  if len(input_B_node.args) > 0:
                      input_B = input_B_node.args[0]
                  else:
                      input_B = input_B_node.kwargs["input"]
  
              if Atrans or Btrans:
                  with module.graph.inserting_before(node):
                      fused_node = module.graph.call_function(
                          transpose_matmul,
                          args=(input_A, input_B, Atrans, Btrans),
                      )
                  node.replace_all_uses_with(fused_node)
                  module.graph.erase_node(node)
                  if Atrans and len(input_A_node.users) == 0:
                      module.graph.erase_node(input_A_node)
                  if Btrans and len(input_B_node.users) == 0:
                      module.graph.erase_node(input_B_node)
  
      module.graph.lint()
      module.recompile()
      return module
  
imports: import copy
  import logging
  from typing import List, Optional
  import torch
  import torch.nn
  from torch._dynamo.utils import detect_fake_mode
  from torch._utils_internal import print_graph
  from torch.fx.experimental.optimization import matches_module_pattern, replace_node_module
  from torch.fx.passes.shape_prop import ShapeProp
  from torch.nn import functional
  from torch.nn.utils.fusion import fuse_conv_bn_eval, fuse_conv_bn_weights
  from ..None import config
  from ..fx_utils import matches_module_function_pattern
  from ..pattern_matcher import init_once_fakemode, PatternMatcherPass, stable_topological_sort
  from ..utils import is_cpu_device
  from .group_batch_fusion import group_batch_fusion_passes
  from .misc_patterns import numpy_compat_normalization
label_desc: Use image node_id 12 for calling a global function with example usage: permute_matmul_fusion(module) and returns: module