content: import inspect
  import json
  import logging
  import os
  import re
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  
  import torch
  import torch.fx
  from torch.fx.graph_module import GraphModule
  from torch.fx.immutable_collections import immutable_dict, immutable_list
  from torch.fx.node import _get_qualified_name
  from torch.fx.passes import graph_drawer
  from torch.fx.passes.shape_prop import TensorMetadata
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  
  def get_target_from_module(mod: torch.nn.Module, target: str):
      """
      Gets `target` from `mod` and returns it. If `target` is empty then returns `mod.`
      """
      if target == "":
          return mod
  
      target_atoms = target.split(".")
      curr_obj = mod
      for i, atom in enumerate(target_atoms):
          if not hasattr(curr_obj, atom):
              raise RuntimeError(
                  f"Node referenced nonexistent target '{'.'.join(target_atoms[:i])}'; "
                  f" original whole target: '{target}'"
              )
          curr_obj = getattr(curr_obj, atom)
      return curr_obj
  
  
  def get_attr(node: torch.fx.Node) -> Any:
      """
      Returns the underlying attr for a given node which
      must be of type get_attr.
      """
      assert node.op == "get_attr", "Expected a get_attr node"
      return get_target_from_module(node.graph.owning_module, str(node.target))
  
  
  def is_acc_op(node_or_target: Union[Callable, torch.fx.Node]) -> bool:
      """
      Returns whether `node_or_target` is an acc_op. If it's a node, then checks whether
      it's a call_function target is from the acc_ops module. Otherwise it's already
      the target, which is similarly checked to see if it's from the acc_ops module.
      """
      if isinstance(node_or_target, torch.fx.Node):
          # All acc_ops are call_functions.
          if node_or_target.op != "call_function":
              return False
          target = node_or_target.target
      else:
          target = node_or_target
      return "acc_ops" in target.__module__
  
  
  def is_acc_op_with_kwarg(
      node_or_target: Union[Callable, torch.fx.Node], kwarg: str
  ) -> bool:
      """
      Helper that inspects `node_or_target` and returns whether it is an acc_op node
      (or a target for an acc_op) that has an arg signature that includes `kwarg`.
      """
      if not is_acc_op(node_or_target):
          return False
  
      target = (
          node_or_target.target
          if isinstance(node_or_target, torch.fx.Node)
          else node_or_target
      )
      assert not isinstance(target, str)
      return kwarg in inspect.signature(inspect.unwrap(target)).parameters
  
  
  def build_raw_tensor_meta(
      shape=None,
      dtype=None,
      requires_grad=None,
      stride=None,
      memory_format=None,
      is_quantized=None,
      qparams=None,
  ):
      return TensorMetadata(**locals())
  
  
  def draw_graph(traced: torch.fx.GraphModule, fname: str, figname: str = "fx_graph"):
      base, ext = os.path.splitext(fname)
      if not ext:
          ext = ".svg"
      _LOGGER.info(f"Writing FX graph to file: {base}{ext}")
      g = graph_drawer.FxGraphDrawer(traced, figname)
      x = g.get_main_dot_graph()
      try:
          getattr(x, "write_" + ext.lstrip("."))(fname)
      except OSError as e:
          _LOGGER.error(f"Failed to write the FX graph due to: {e}")
  
  
  def get_model_info_str(gm: torch.fx.GraphModule, header: Optional[str] = None):
      """
      Print out info of the provided `gm`.
      If `header` is provided then it's included in the printed string.
      """
      ops_and_counts: Dict[Callable, int] = {}
      placeholder_count = get_attr_count = 0
      call_method_count = call_module_count = output_count = 0
      for node in gm.graph.nodes:
          if node.op == "call_function":
              ops_and_counts[node.target] = ops_and_counts.get(node.target, 0) + 1
          elif node.op == "placeholder":
              placeholder_count += 1
          elif node.op == "get_attr":
              get_attr_count += 1
          elif node.op == "call_method":
              call_method_count += 1
          elif node.op == "call_module":
              call_module_count += 1
          elif node.op == "output":
              output_count = len(node.args[0]) if isinstance(node.args[0], tuple) else 1
          else:
              raise RuntimeError(f"Unknown node found: {node.format_node()}")
  
      header = "" if header is None else f" [{header}]"
      model_info_str = f"Model Info{header}:\n"
      model_info_str += f"> placeholder: {placeholder_count}\n"
      model_info_str += f"> get_attr: {get_attr_count}\n"
      model_info_str += f"> output: {output_count}\n"
      if call_module_count != 0:
          model_info_str += f"> WARNING: call_module: {call_module_count}"
      if call_method_count != 0:
          model_info_str += f"> WARNING: call_method: {call_method_count}"
  
      # Sort and print all the other ops. Sort so it's deterministic between runs and
      # easier to parse.
      pretty_ops_and_counts: List[Tuple[str, int]] = []
      for op, count in ops_and_counts.items():
          name = strip_module_prefixes(_get_qualified_name(op))
          pretty_ops_and_counts.append((name, count))
      pretty_ops_and_counts.sort()
      for op_str, count in pretty_ops_and_counts:
          model_info_str += f"> {op_str}: {count}\n"
  
      return model_info_str
  
  
  def strip_module_prefixes(op_name):
      return (
          op_name.replace("torch_tensorrt.fx.tracer.acc_tracer.", "")
          .replace("glow.fb.fx.acc_tracer.", "")
          .replace("glow.fb.fx.", "")
      )
  
  
  def get_unique_attr_name_in_module(mod_traced: torch.fx.GraphModule, name: str) -> str:
      """
      Make sure the name is unique (in a module) and can represents an attr.
      """
      # Delete all characters that are illegal in a Python identifier.
      name = re.sub("[^0-9a-zA-Z_]+", "_", name)
      if name[0].isdigit():
          name = f"_{name}"
      # Now make sure it is in fact unique to the module by incrementing suffix value.
      while hasattr(mod_traced, name):
          match = re.match(r"(.*)_(\d+)$", name)
          if match is None:
              name = name + "_1"
          else:
              base, num = match.group(1, 2)
              name = f"{base}_{int(num) + 1}"
  
      return name
  
  
  def map_tensor_metadata(a: Any, fn: Callable):
      """
      Map some `fn` to `a`, where `a` is either a TensorMetadata, or else a tuple/list/dict
      recursively containing TensorMetadata.
      """
      if isinstance(a, int):
          return 1
      elif a is None:
          return 1
      elif isinstance(a, TensorMetadata):
          return fn(a)
      elif isinstance(a, tuple):
          return tuple(map_tensor_metadata(elem, fn) for elem in a)
      elif isinstance(a, dict):
          return immutable_dict(
              {name: map_tensor_metadata(elem, fn) for name, elem in a.items()}
          )
      assert isinstance(
          a, list
      ), f"Only supporting tuple/list/TensorMetadata, but found {type(a)}"
      return immutable_list(map_tensor_metadata(elem, fn) for elem in a)
  
  
  def get_tensor_meta(node: torch.fx.Node) -> TensorMetadata:
      tensor_meta = node.meta.get("tensor_meta")
  
      if not tensor_meta:
          raise RuntimeError(
              f"Node has no tensor metadata associated with it! "
              f"Check that shape propagation has run. {node.format_node()}"
          )
      return tensor_meta
  