code: def test_custom_pre_grad_pass(self):
      saved_graph = [None]
  
      def merge_mm_shared_rhs(graph: fx.Graph):
          """
          Bad POC of merging mm with a shared RHS.
          i.e. [mm(x, W), mm(x2, W)] => mm(cat(x, x2), W).split()
  
          Isn't actually safe for a couple reasons. For example, it doesn't handle the
          case where the LHS inputs depend on each other
          """
          saved_graph[0] = graph
          matmuls = [n for n in graph.nodes if n.target == torch.mm]
          rhs_vals = defaultdict(set)
          for m in matmuls:
              rhs_vals[m.args[1]].add(m)
  
          order = {}
          for idx, n in enumerate(graph.nodes):
              order[n] = idx
  
          for rhs, matmuls in rhs_vals.items():
              if len(matmuls) == 1:
                  continue
              matmuls = sorted(matmuls, key=lambda x: order[x])
              with graph.inserting_before(matmuls[0]):
                  lhs_vals = [m.args[0] for m in matmuls]
                  new_cat = graph.create_node(
                      "call_function", torch.cat, args=(lhs_vals, 0)
                  )
                  new_mm = graph.create_node(
                      "call_function", torch.mm, args=(new_cat, rhs)
                  )
                  split_vals = graph.create_node(
                      "call_function",
                      torch.split,
                      args=(
                          new_mm,
                          [
                              l.meta["example_value"].shape[0]
                              for l in lhs_vals
                          ],
                      ),
                  )
              for idx, m in enumerate(matmuls):
                  m.target = operator.getitem
                  m.args = (split_vals, idx)
  
      @config.patch(pre_grad_custom_pass=merge_mm_shared_rhs)
      def inner_test():
          @torch.compile
          def f(W, nested_seqs):
              outs = [torch.mm(s, W) for s in nested_seqs]
              return outs
  
          W = torch.randn(16, 16, dtype=torch.bfloat16)
          nested_seqs = [
              torch.randn(l, 16, dtype=torch.bfloat16)
              for l in [4, 8, 5, 3]
          ]
  
          f(W, nested_seqs)
          assert saved_graph[0] is not None
          matmuls = [
              n for n in saved_graph[0].nodes if n.target == torch.mm
          ]
          assert len(matmuls) == 1
  
      inner_test()
  
imports: import contextlib
  import operator
  from collections import defaultdict
  import torch
  import torch._inductor.pattern_matcher
  import torch.fx
  from torch._dynamo.test_case import run_tests, TestCase
  from torch._dynamo.utils import counters
  from torch._inductor import config
  from torch._inductor.lowering import lowerings
  from torch._inductor.pattern_matcher import Arg, CallFunction, PatternMatcherPass
  from torch.testing._internal.common_utils import IS_LINUX
  from torch.testing._internal.inductor_utils import HAS_CPU
label_desc: Use image node_id 4 for calling the TestPostGradCustomPrePostPass obj's underlying member method code with example usage: obj.test_custom_pre_grad_pass() and returns: outs