class_docstr: Generates model graph, each node is created from single or multiple jit trace nodes.
class_docstr_tok: ["Generates","model","graph",",","each","node","is","created","from","single","or","multiple","jit","trace","nodes","."]
docstr: For trace graph nodes, some nodes are not in modules, these nodes are usually generated by
  the functions directly called in module ```forward```. For such nodes, some of them are
  trivial op which are label by ```prim::```, some of them are not such ops which is call
  non-prim ops. This function is to merge neighbor prim ops to a non-prim op, to construct
  a node.
  
  Parameters
  ----------
  node : trace graph node
      The non-prim node to expand
  nodes : list of trace graph node
      All the trace graph nodes within the same scope as the non-prim node
  input_to_node : dict
      key: input name, value: a node that uses this input
  output_to_node : dict
      key: output name, value: a node that generates this output
  module_type : str
      can be 'module' or 'func'
  
  Returns
  -------
  node
      the expanded non-prim node
docstr_tok: ["For","trace","graph","nodes",",","some","nodes","are","not","in","modules",",","these","nodes","are","usually","generated","by","the","functions","directly","called","in","module","``","`","forward","``","`",".","For","such","nodes",",","some","of","them","are","trivial","op","which","are","label","by","``","`","prim",":",":","``","`",",","some","of","them","are","not","such","ops","which","is","call","non-prim","ops",".","This","function","is","to","merge","neighbor","prim","ops","to","a","non-prim","op",",","to","construct","a","node",".","Parameters","--","--","--","--","--","node",":","trace","graph","node","The","non-prim","node","to","expand","nodes",":","list","of","trace","graph","node","All","the","trace","graph","nodes","within","the","same","scope","as","the","non-prim","node","input_to_node",":","dict","key",":","input","name",",","value",":","a","node","that","uses","this","input","output_to_node",":","dict","key",":","output","name",",","value",":","a","node","that","generates","this","output","module_type",":","str","can","be","'module","'","or","'func'","Returns","--","--","--","-","node","the","expanded","non-prim","node"]
code: def _expand_key_func_node(
      self, node, nodes, input_to_node, output_to_node, module_type
  ):
      """
      For trace graph nodes, some nodes are not in modules, these nodes are usually generated by
      the functions directly called in module ```forward```. For such nodes, some of them are
      trivial op which are label by ```prim::```, some of them are not such ops which is call
      non-prim ops. This function is to merge neighbor prim ops to a non-prim op, to construct
      a node.
  
      Parameters
      ----------
      node : trace graph node
          The non-prim node to expand
      nodes : list of trace graph node
          All the trace graph nodes within the same scope as the non-prim node
      input_to_node : dict
          key: input name, value: a node that uses this input
      output_to_node : dict
          key: output name, value: a node that generates this output
      module_type : str
          can be 'module' or 'func'
  
      Returns
      -------
      node
          the expanded non-prim node
      """
      # TODO: scope name could be empty
      node_name = ".".join(
          [
              self._get_module_name(node.scopeName()),
              node.kind(),
              str(self.global_count),
          ]
      )
      unique_name = node_name
      _logger.debug("expand non-prim node, node name: %s", node_name)
      self.global_count += 1
      op_type = node.kind()
      node_group = [node]
      inputs = []
      outputs = []
      node_queue = queue.Queue()
      node_queue.put(node)
      while not node_queue.empty():
          curr_node = node_queue.get()
          for _input in curr_node.inputs():
              if _input.node().kind() == CONSTANT_KIND:
                  continue
              input_name = _input.debugName()
              if input_name in output_to_node:
                  for predecessor_node in output_to_node[input_name]:
                      if predecessor_node in nodes:
                          if not self._is_key_func(predecessor_node):
                              if predecessor_node not in node_group:
                                  node_group.append(predecessor_node)
                                  node_queue.put(predecessor_node)
                          else:
                              inputs.append(input_name)
                      else:
                          inputs.append(input_name)
              else:
                  inputs.append(input_name)
      for output in node.outputs():
          if output.node().kind() == CONSTANT_KIND:
              continue
          outputs.append(output.debugName())
      nodepy = NodePyGroup(
          node_name,
          unique_name,
          module_type,
          op_type,
          node_group,
          inputs=inputs,
          outputs=outputs,
          key_node=node,
      )
      return nodepy
  
imports: import logging
  import queue
  import re
  from collections import defaultdict
  from typing import List, Dict
  import torch
  from torch.utils.tensorboard._pytorch_graph import NodePy, NodePyIO, NodePyOP, GraphPy
label_desc: Use image node_id 2 for calling the TorchModuleGraph obj's underlying member method code with example usage: obj._expand_key_func_node(node, nodes, input_to_node, output_to_node, module_type) and returns: nodepy