content: import logging
  import math
  import operator
  import os
  from dataclasses import dataclass, field
  from typing import Any, Dict, List, Union
  
  import torch
  from torch_tensorrt.dynamo._settings import CompilationSettings
  from torch_tensorrt.dynamo.conversion._ConverterRegistry import ConverterRegistry
  from torch_tensorrt.dynamo.conversion.converter_utils import get_node_name
  
  logger = logging.getLogger(__name__)
  
  
  @dataclass
  class PerSubgraphData:
      """Class to track data on a per-subgraph level
  
      Args:
          subgraph_name (str): Name of the subgraph in the GraphModule
          subgraph_op_count (int): Number of operations in the subgraph
          input_shapes (Any): Shapes of input Tensors of the subgraph
          input_dtypes (Any): Input data types of the subgraph
          output_shapes (Any): Shapes of output Tensors of the subgraph
          output_dtypes (Any): Output data types of the subgraph
      """
  
      subgraph_name: str = ""
      subgraph_op_count: int = 0
      input_shapes: Any = field(default_factory=list)
      input_dtypes: Any = field(default_factory=list)
      output_shapes: Any = field(default_factory=list)
      output_dtypes: Any = field(default_factory=list)
  
  
  @dataclass
  class DryRunTracker:
      """Class to track data on a graph-wide level
  
      Args:
          total_ops_in_graph (int): Total number of operators in graph
          supported_ops_in_graph (int): Number of supported operators in graph
          input_shapes (Any): Shapes of input Tensors of the graph
          input_dtypes (Any): Input data types of the graph
          output_shapes (Any): Shapes of output Tensors of the graph
          output_dtypes (Any): Output data types of the graph
          per_subgraph_data (List[PerSubgraphData]): Per-subgraph data, see above class
          tensorrt_graph_count (int): Number of TensorRT engines to be generated
          compilation_settings (CompilationSettings): User Compilation Settings
          unsupported_ops (Dict[str, int]): Set of operators not supported in TRT
          to_run_in_torch (List[str]): Set of nodes to run in Torch
      """
  
      total_ops_in_graph: int = 0
      supported_ops_in_graph: int = 0
      input_shapes: Any = field(default_factory=list)
      input_dtypes: Any = field(default_factory=list)
      output_shapes: Any = field(default_factory=list)
      output_dtypes: Any = field(default_factory=list)
      per_subgraph_data: List[PerSubgraphData] = field(default_factory=list)
      tensorrt_graph_count: int = 0
      compilation_settings: CompilationSettings = field(
          default_factory=CompilationSettings
      )
      unsupported_ops: Dict[str, int] = field(default_factory=dict)
      to_run_in_torch: List[str] = field(default_factory=list)
  
  
  def dryrun_stats_display(
      dryrun_tracker: DryRunTracker, dryrun_enabled: Union[bool, str]
  ) -> None:
      """Displays statistics about the dryrun either to debug logs or stdout"""
      formatted_stats = "\n"
  
      # Print overall stats about the graph, operator counts, etc.
      formatted_stats += "+" * 50 + " Dry-Run Results for Graph " + "+" * 50 + "\n\n"
      formatted_stats += (
          f"The graph consists of {dryrun_tracker.total_ops_in_graph} Total Operators, "
          f"of which {dryrun_tracker.supported_ops_in_graph} operators are supported, "
          f"{round(dryrun_tracker.supported_ops_in_graph*100/dryrun_tracker.total_ops_in_graph, 2)}% coverage\n\n"
      )
      if dryrun_tracker.unsupported_ops:
          parsed_ops = "\n".join(
              [f"{str(k)}: {str(v)}" for k, v in dryrun_tracker.unsupported_ops.items()]
          )
          formatted_stats += f"The following ops are currently unsupported or excluded from conversion, and are listed with their op-count in the graph:\n {parsed_ops}\n\n"
  
      if dryrun_tracker.to_run_in_torch:
          formatted_nodes = "\n".join(dryrun_tracker.to_run_in_torch)
          formatted_stats += (
              f"The following nodes are currently set to run in Torch:\n{formatted_nodes}\n"
              "Note: Some of the above nodes may be supported, but were not included in a TRT graph by the partitioner\n\n"
          )
  
      formatted_stats += f"Compiled with: {dryrun_tracker.compilation_settings}\n\n"
  
      assert len(dryrun_tracker.per_subgraph_data) == dryrun_tracker.tensorrt_graph_count
  
      # Print schematic of the graph structure, as in:
      #
      #   Inputs: List[Tensor: (1, 3, 224, 224)@float32]
      #    ...
      #      TRT Engine #1 - Submodule name: _run_on_acc_0
      #       Engine Inputs: List[Tensor: (1, 3, 224, 224)@float32]
      #       Number of Operators in Engine: 1
      #       Engine Outputs: Tensor: (1, 64, 112, 112)@float32
      #    ...
      #   Outputs: List[Tensor: (1, 1000)@float32]
      #
      formatted_stats += " " * 2 + "Graph Structure:\n\n"
      formatted_stats += (
          " " * 3
          + f"Inputs: {input_formatter(dryrun_tracker.input_shapes, dryrun_tracker.input_dtypes)}\n"
      )
  
      for i, trt_subgraph_data in enumerate(dryrun_tracker.per_subgraph_data):
          formatted_stats += " " * 4 + "...\n"
          formatted_stats += (
              " " * 4
              + f"TRT Engine #{i+1} - Submodule name: {trt_subgraph_data.subgraph_name}\n"
          )
          formatted_stats += (
              " " * 5
              + f"Engine Inputs: {input_formatter(trt_subgraph_data.input_shapes, trt_subgraph_data.input_dtypes)}\n"
          )
          formatted_stats += (
              " " * 5
              + f"Number of Operators in Engine: {trt_subgraph_data.subgraph_op_count}\n"
          )
          formatted_stats += (
              " " * 5
              + f"Engine Outputs: {input_formatter(trt_subgraph_data.output_shapes, trt_subgraph_data.output_dtypes)}\n"
          )
  
      formatted_stats += " " * 4 + "...\n"
      formatted_stats += (
          " " * 3
          + f"Outputs: {input_formatter(dryrun_tracker.output_shapes, dryrun_tracker.output_dtypes)}\n"
      )
  
      # Print aggregate statistics about the graph structure, including recommended "min_block_size" options
      if dryrun_tracker.tensorrt_graph_count > 0:
          min_ops_in_an_engine = min(
              trt_subgraph.subgraph_op_count
              for trt_subgraph in dryrun_tracker.per_subgraph_data
          )
          avg_ops_per_engine = (
              sum(
                  trt_subgraph.subgraph_op_count
                  for trt_subgraph in dryrun_tracker.per_subgraph_data
              )
              / dryrun_tracker.tensorrt_graph_count
          )
          avg_ops_per_engine = round(avg_ops_per_engine, 2)
          most_ops_in_an_engine = max(
              trt_subgraph.subgraph_op_count
              for trt_subgraph in dryrun_tracker.per_subgraph_data
          )
  
          formatted_stats += "\n" + " " * 2 + "-" * 25 + " Aggregate Stats " + "-" * 25
          formatted_stats += (
              "\n\n"
              + " " * 3
              + "Average Number of Operators per TRT Engine: "
              + f"{avg_ops_per_engine}"
          )
  
          formatted_stats += (
              "\n"
              + " " * 3
              + "Most Operators in a TRT Engine: "
              + f"{most_ops_in_an_engine}"
          )
  
          formatted_stats += "\n\n" + " " * 2 + "*" * 10 + " Recommendations " + "*" * 10
          formatted_stats += (
              "\n\n"
              + " " * 3
              + "- For minimal graph segmentation, select min_block_size="
              + f"{most_ops_in_an_engine} which would generate "
              + f"{len([1 for trt_subgraph in dryrun_tracker.per_subgraph_data if trt_subgraph.subgraph_op_count >= most_ops_in_an_engine])} TRT engine(s)"
          )
          if math.ceil(avg_ops_per_engine) != most_ops_in_an_engine:
              formatted_stats += (
                  "\n"
                  + " " * 3
                  + "- For moderate graph segmentation, select min_block_size="
                  + f"{math.ceil(avg_ops_per_engine)} which would generate "
                  + f"{len([1 for trt_subgraph in dryrun_tracker.per_subgraph_data if trt_subgraph.subgraph_op_count >= math.ceil(avg_ops_per_engine)])} TRT engine(s)"
              )
  
          formatted_stats += (
              "\n"
              + " " * 3
              + "- The current level of graph segmentation is equivalent to selecting min_block_size="
              + f"{min_ops_in_an_engine} which generates "
              + f"{len([1 for trt_subgraph in dryrun_tracker.per_subgraph_data if trt_subgraph.subgraph_op_count >= min_ops_in_an_engine])} TRT engine(s)"
          )
      else:
          formatted_stats += (
              "\n"
              + " " * 2
              + "Aggregate stats not available since no TRT Engines were generated."
          )
  
      # If user specified "dryrun=True", print to stdout, else debug
      # If user specified a filepath, save the output to the path as well
      if dryrun_enabled:
          print(formatted_stats)
          if isinstance(dryrun_enabled, str):
              if os.path.exists(dryrun_enabled):
                  logger.warning(
                      f"File already exists at path {dryrun_enabled}, not saving dryrun output"
                  )
              else:
                  with open(dryrun_enabled, "w+") as f:
                      f.write(formatted_stats)
      else:
          logger.debug(formatted_stats)
  
  
  def input_formatter(shapes: Any, dtypes: Any) -> str:
      """Format shapes and dtypes of input Tensors into a readable string"""
  
      def input_formatter_helper(shapes: Any, dtypes: Any) -> str:
          """Helper for input formatter"""
          # Base case 1 - single static/dynamic shape, single dtype
          if isinstance(shapes, tuple) and all(
              isinstance(elt, (int, tuple)) for elt in shapes
          ):
              input_shape_string = "Tensor: ("
              for elt in shapes:
                  if isinstance(elt, tuple):
                      input_shape_string += f"(min={elt[0]}, max={elt[1]}), "
                  else:
                      input_shape_string += f"{elt}, "
              input_shape_string = input_shape_string[:-2] + ")" + f"@{str(dtypes)[6:]}, "
              return input_shape_string
  
          # Base case 2 - dynamic shape, single dtype
          elif (
              isinstance(shapes, dict)
              and len(shapes) == 3
              and all(
                  (
                      isinstance(shape, tuple)
                      and all(isinstance(elt, int) for elt in shape)
                      and k in ("min_shape", "opt_shape", "max_shape")
                  )
                  for k, shape in shapes.items()
              )
          ):
              return f"Tensor: {shapes}@{str(dtypes)[6:]}, "
  
          # Shapes is a sequence
          elif isinstance(shapes, (list, tuple)):
              formatted_str = "List[" if isinstance(shapes, list) else "Tuple("
              for shape, dtype in zip(shapes, dtypes):
                  formatted_str += input_formatter_helper(shape, dtype)
              formatted_str = formatted_str[:-2] + (
                  "], " if isinstance(shapes, list) else "), "
              )
              return formatted_str
  
          # Shapes is a dictionary
          elif isinstance(shapes, dict):
              formatted_str = "Dict{"
  
              for key, shape in shapes.items():
                  formatted_str += input_formatter_helper(shape, dtypes[key])
  
              formatted_str = formatted_str[:-2] + "}, "
              return formatted_str
  
          else:
              raise ValueError(
                  f"Invalid input type {type(shapes)} encountered in parse_complex_tensor_structs parsing."
              )
  
      return input_formatter_helper(shapes, dtypes)[:-2]
  
  
  def parse_non_trt_nodes(graph_module: torch.fx.GraphModule) -> List[str]:
      """Parses call_function and call_method nodes from a GraphModule
      Excludes getitem nodes
  
      Returns a string representation of the nodes
      """
      to_run_in_torch = []
      for node in graph_module.graph.nodes:
          # getitem nodes are excluded since they are a Tensor-collection op
          if (
              node.op in ("call_function", "call_method")
              and node.target != operator.getitem
          ):
              to_run_in_torch.append(
                  f"Node: {ConverterRegistry.qualified_name_or_str(node.target)}, "
                  f"with layer location: {get_node_name(node)}"
              )
      return to_run_in_torch
  