code: def test_split_by_tags(self) -> None:
      tags = ["red", "blue", "green"]
      module = TestSplitByTags.TestModule()
      gm, tag_node = TestSplitByTags.trace_and_tag(module, tags)
      split_gm, orig_to_split_fqn_mapping = split_by_tags(
          gm, tags, return_fqn_mapping=True
      )
      # Ensure split_gm has (and only has) ordered submodules named
      # red_0, blue_1, green_2
      for idx, (name, _) in enumerate(split_gm.named_children()):
          if idx < len(tags):
              self.assertTrue(
                  name == tags[idx],
                  f"split_gm has an incorrect submodule named {name}",
              )
  
      # Ensure each submodule has expected (ordered) call_module node(s).
      # For example, a submodule named split_gm.red_0 has (and only has) linear1;
      # split_gm.green_2 has (and only has) linear3 and linear4 with order
      sub_graph_idx = 0
      for sub_name, sub_graph_module in split_gm.named_children():
          node_idx = 0
          for node in sub_graph_module.graph.nodes:
              if node.op != "call_module":
                  continue
              self.assertTrue(
                  node.name == tag_node[f"{sub_name}"][node_idx],
                  # pyre-fixme[61]: `name` is undefined, or not always defined.
                  f"{sub_name} has incorrectly include {node.name}",
              )
              node_idx += 1
          sub_graph_idx += 1
  
      self.assertEqual(
          orig_to_split_fqn_mapping,
          {
              "linear1": "red.linear1",
              "linear2": "blue.linear2",
              "linear3": "green.linear3",
              "linear4": "green.linear4",
          },
          f"{orig_to_split_fqn_mapping=}",
      )
  
file_path: pytorch/test/fx/test_fx_split.py
imports: from collections import defaultdict
  from typing import List, Tuple, Dict
  import torch
  from torch.fx.passes.split_utils import split_by_tags
  from torch.testing._internal.common_utils import TestCase
label_desc: Use image node_id 2 for calling the TestSplitByTags obj's underlying member method code with example usage: obj.test_split_by_tags() without return types