docstr: Test simple gm consists of nodes with tag (only show call_module nodes here):
      linear1 - tag: "red"
      linear2 - tag: "blue"
      linear3, linear4 - tag: "green"
  
  At the beginning we have:
      gm:
          linear1
          linear2
          linear3
          linear4
  
  split_gm = split_by_tags(gm, tags)
  
  Then we have:
      split_gm:
          red:
              linear1
          blue:
              linear2
          green:
              linear3
              linear4
docstr_tok: ["Test","simple","gm","consists","of","nodes","with","tag","(","only","show","call_module","nodes","here",")",":","linear1","-","tag",":","``","red","''","linear2","-","tag",":","``","blue","''","linear3",",","linear4","-","tag",":","``","green","''","At","the","beginning","we","have",":","gm",":","linear1","linear2","linear3","linear4","split_gm","=","split_by_tags","(","gm",",","tags",")","Then","we","have",":","split_gm",":","red",":","linear1","blue",":","linear2","green",":","linear3","linear4"]
code: def trace_and_tag(
      module: torch.nn.Module, tags: List[str]
  ) -> Tuple[torch.fx.GraphModule, Dict[str, List[str]]]:
      """
      Test simple gm consists of nodes with tag (only show call_module nodes here):
          linear1 - tag: "red"
          linear2 - tag: "blue"
          linear3, linear4 - tag: "green"
  
      At the beginning we have:
          gm:
              linear1
              linear2
              linear3
              linear4
  
      split_gm = split_by_tags(gm, tags)
  
      Then we have:
          split_gm:
              red:
                  linear1
              blue:
                  linear2
              green:
                  linear3
                  linear4
      """
      tag_node = defaultdict(list)
      gm: torch.fx.GraphModule = torch.fx.symbolic_trace(module)
  
      # Add tag to all nodes and build dictionary record tag to call_module nodes
      for node in gm.graph.nodes:
          if "linear1" in node.name:
              node.tag = tags[0]
              tag_node[tags[0]].append(node.name)
          elif "linear2" in node.name:
              node.tag = tags[1]
              tag_node[tags[1]].append(node.name)
          else:
              node.tag = tags[2]
              if node.op == "call_module":
                  tag_node[tags[2]].append(node.name)
      return gm, tag_node
  
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 1 for calling the TestSplitByTags obj's underlying member method code with example usage: obj.trace_and_tag(module, tags) and returns: gm,  tag_node