class_docstr: 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
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 = set()
      outputs = set()
      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.add(input_name)
                      else:
                          inputs.add(input_name)
              else:
                  inputs.add(input_name)
      for output in node.outputs():
          if output.node().kind() == CONSTANT_KIND:
              continue
          outputs.add(output.debugName())
      nodepy = NodePyGroup(
          node_name,
          unique_name,
          module_type,
          op_type,
          node_group,
          inputs=list(inputs),
          outputs=list(outputs),
          key_node=node,
      )
      return nodepy
  
imports: import logging
  import queue
  import re
  from collections import defaultdict
  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