code: def _get_count(param_count: Dict, node_name: str) -> int:
      """Identify different mutations of a given node name."""
      # TODO(anj): This is not very stable since it is possible that the name
      # may not be in the same format. Is there another way to identify nodes
      # in a graph?
      if node_name in param_count:
          return param_count[node_name]
      elif node_name.split("_")[0] in param_count:
          return param_count[node_name.split("_")[0]]
      else:
          raise RuntimeError(
              f"Unable to find match between param {param_count} and node {node_name}"
          )
  
file_path: fairscale/fairscale/experimental/nn/auto_shard.py
imports: import logging
  from typing import Dict, List, Set
  import torch
  import torch.fx
  from torch.fx.node import Node
label_desc: Use image node_id 1 for calling a global function with example usage: _get_count(param_count, node_name) and returns: param_count, param_count