content: from __future__ import annotations
  
  import ctypes
  import gc
  import logging
  import os
  import platform
  import warnings
  from dataclasses import fields, replace
  from enum import Enum
  from typing import (
      Any,
      Callable,
      Dict,
      List,
      Optional,
      Sequence,
      Tuple,
      Union,
  )
  
  import numpy as np
  import psutil
  import sympy
  import tensorrt as trt
  import torch
  from torch._subclasses.fake_tensor import FakeTensor
  from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
  from torch_tensorrt._Device import Device
  from torch_tensorrt._enums import dtype
  from torch_tensorrt._features import ENABLED_FEATURES
  from torch_tensorrt._Input import Input
  from torch_tensorrt._utils import is_tensorrt_version_supported
  from torch_tensorrt.dynamo import _defaults
  from torch_tensorrt.dynamo._defaults import default_device
  from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
  from torch_tensorrt.dynamo._settings import CompilationSettings
  
  from packaging import version
  
  from .types import TRTDataType
  
  logger = logging.getLogger(__name__)
  
  COSINE_THRESHOLD = 0.99
  DYNAMIC_DIM = -1
  RTOL = 5e-3
  ATOL = 5e-3
  CPU_DEVICE = "cpu"
  
  
  class Frameworks(Enum):
      NUMPY = "numpy"
      TORCH = "torch"
      TRT = "trt"
  
  
  DataTypeEquivalence: Dict[
      TRTDataType, Dict[Frameworks, Union[TRTDataType, np.dtype, torch.dtype]]
  ] = {
      trt.int8: {
          Frameworks.NUMPY: np.int8,
          Frameworks.TORCH: torch.int8,
          Frameworks.TRT: trt.int8,
      },
      trt.int32: {
          Frameworks.NUMPY: np.int32,
          Frameworks.TORCH: torch.int32,
          Frameworks.TRT: trt.int32,
      },
      trt.int64: {
          Frameworks.NUMPY: np.int64,
          Frameworks.TORCH: torch.int64,
          Frameworks.TRT: trt.int64,
      },
      trt.float16: {
          Frameworks.NUMPY: np.float16,
          Frameworks.TORCH: torch.float16,
          Frameworks.TRT: trt.float16,
      },
      trt.float32: {
          Frameworks.NUMPY: np.float32,
          Frameworks.TORCH: torch.float32,
          Frameworks.TRT: trt.float32,
      },
      trt.bool: {
          Frameworks.NUMPY: bool,
          Frameworks.TORCH: torch.bool,
          Frameworks.TRT: trt.bool,
      },
  }
  
  if is_tensorrt_version_supported("7.0"):
      DataTypeEquivalence[trt.bool] = {
          Frameworks.NUMPY: np.bool_,
          Frameworks.TORCH: torch.bool,
          Frameworks.TRT: trt.bool,
      }
  
  
  def unified_dtype_converter(
      dtype: Union[TRTDataType, torch.dtype, np.dtype], to: Frameworks
  ) -> Union[np.dtype, torch.dtype, TRTDataType]:
      """
      Convert TensorRT, Numpy, or Torch data types to any other of those data types.
      Args:
          dtype (TRTDataType, torch.dtype, np.dtype): A TensorRT, Numpy, or Torch data type.
          to (Frameworks): The framework to convert the data type to.
      Returns:
          The equivalent data type in the requested framework.
      """
      assert to in Frameworks, f"Expected valid Framework for translation, got {to}"
      trt_major_version = int(trt.__version__.split(".")[0])
      if dtype in (np.int8, torch.int8, trt.int8):
          return DataTypeEquivalence[trt.int8][to]
      elif trt_major_version >= 7 and dtype in (np.bool_, torch.bool, trt.bool):
          return DataTypeEquivalence[trt.bool][to]
      elif dtype in (np.int32, torch.int32, trt.int32):
          return DataTypeEquivalence[trt.int32][to]
      elif dtype in (np.int64, torch.int64, trt.int64):
          return DataTypeEquivalence[trt.int64][to]
      elif dtype in (np.float16, torch.float16, trt.float16):
          return DataTypeEquivalence[trt.float16][to]
      elif dtype in (np.float32, torch.float32, trt.float32):
          return DataTypeEquivalence[trt.float32][to]
      else:
          raise TypeError("%s is not a supported dtype" % dtype)
  
  
  def deallocate_module(module: torch.fx.GraphModule, delete_module: bool = True) -> None:
      """
      This is a helper function to delete the instance of module. We first move it to CPU and then
      delete the object. This function ensures the GPU memory occupied by the module is released effectively after this call
      """
      module.to(CPU_DEVICE)
      if delete_module:
          del module
      torch.cuda.empty_cache()
      gc.collect()
  
  
  def use_python_runtime_parser(use_python_runtime: Optional[bool] = None) -> bool:
      """Parses a user-provided input argument regarding Python runtime
  
      Automatically handles cases where the user has not specified a runtime (None)
  
      Returns True if the Python runtime should be used, False if the C++ runtime should be used
      """
      using_python_runtime = use_python_runtime
      reason = ""
  
      # Runtime was manually specified by the user
      if using_python_runtime is not None:
          reason = "as requested by user"
      # Runtime was not manually specified by the user, automatically detect runtime
      else:
          try:
              from torch_tensorrt.dynamo.runtime import TorchTensorRTModule  # noqa: F401
  
              using_python_runtime = False
              reason = "since C++ dependency was detected as present"
          except ImportError:
              using_python_runtime = True
              reason = "since import failed, C++ dependency not installed"
  
      logger.info(
          f"Using {'Python-only' if using_python_runtime else 'Default'} Torch-TRT Runtime ({reason})"
      )
  
      return using_python_runtime
  
  
  def cosine_similarity(gt_tensor: torch.Tensor, pred_tensor: torch.Tensor) -> float:
      gt_tensor = gt_tensor.flatten().to(torch.float32)
      pred_tensor = pred_tensor.flatten().to(torch.float32)
      if torch.sum(gt_tensor) == 0.0 or torch.sum(pred_tensor) == 0.0:
          if torch.allclose(gt_tensor, pred_tensor, atol=1e-4, rtol=1e-4, equal_nan=True):
              return 1.0
      res_t = torch.nn.functional.cosine_similarity(
          gt_tensor, pred_tensor, dim=0, eps=1e-6
      )
      res: float = res_t.cpu().detach().item()
  
      return res
  
  
  def input_is_dynamic(inputs: Sequence[Union[Input, torch.Tensor, FakeTensor]]) -> bool:
      """
      Return true if any of inputs have dynamic shapes. Supported types are torch_tensorrt.Input | torch.Tensor
      """
      for input in inputs:
          if isinstance(input, torch.Tensor):
              if contains_sym_int(input.shape):
                  return True
          elif isinstance(input, Input):
              if input.shape_mode == Input._ShapeMode.DYNAMIC:
                  return True
          else:
              raise AssertionError(
                  f"Invalid input type ({type(input)}) found. Supported types are torch_tensorrt.Input | torch.Tensor"
              )
  
      return False
  
  
  def get_torch_tensor(
      input: Input,
      device: torch.device,
      mode: str = "",
  ) -> Union[int, torch.Tensor]:
      if input.is_shape_tensor:
          # TODO: All the shape tensors we've encountered so far are plain integers.
          # Validate this assumption on more models.
          return input.shape["opt_shape"][0]
  
      if len(mode) > 0:
          return input.example_tensor(mode).to(device)
      else:
          return input.torch_tensor.to(device)
  
  
  def get_torch_inputs(
      inputs: Sequence[Input] | Dict[str, Any],
      device: Union[Device, torch.device, str],
      mode: str = "",
  ) -> Sequence[Union[int, torch.Tensor]] | Dict[str, Union[int, torch.Tensor]]:
      """
      Return the torch_tensor from the Input object. If mode is set, this implies
      user is using dynamic shaped inputs and return the corresponding input based
      on the mode requested.
      """
      device = to_torch_device(device)
  
      if isinstance(inputs, dict):
          result_dict: Dict[str, Union[int, torch.Tensor]] = {}
          for k, v in inputs.items():
              if isinstance(v, (list, tuple, dict)):
                  result_dict[k] = get_torch_inputs(v, device)
              elif isinstance(v, Input):
                  result_dict[k] = get_torch_tensor(v, device, mode)
          return result_dict
      else:
          result_list: List[Union[int, torch.Tensor]] = []
          for input in inputs:
              if isinstance(input, Input):
                  result_list.append(get_torch_tensor(input, device, mode))
              elif isinstance(input, torch.Tensor):
                  result_list.append(input.to(device))
              else:
                  raise AssertionError(f"Input type {type(input)} is not a valid type")
          return result_list
  
  
  def get_model_device(module: torch.fx.GraphModule) -> torch.device:
      """
      Returns the device on which the module parameters exist.
      """
      device = None
      for parameter in list(module.parameters()):
          if isinstance(parameter, (torch.nn.parameter.Parameter, torch.Tensor)):
              return parameter.device
  
      for buffer in list(module.buffers()):
          if isinstance(buffer, (torch.Tensor)):
              return buffer.device
  
      if device is None:
          device = to_torch_device(default_device())
  
      return device
  
  
  def set_log_level(parent_logger: Any, level: Any) -> None:
      """
      Sets the log level to the user provided level.
      This is used to set debug logging at a global level
      at entry points of tracing, dynamo and torch_compile compilation.
      And set log level for c++ torch trt logger if runtime is available.
      """
      if parent_logger:
          parent_logger.setLevel(level)
  
      if ENABLED_FEATURES.torch_tensorrt_runtime:
          if level == logging.DEBUG:
              log_level = trt.ILogger.Severity.VERBOSE
          elif level == logging.INFO:
              log_level = trt.ILogger.Severity.INFO
          elif level == logging.WARNING:
              log_level = trt.ILogger.Severity.WARNING
          elif level == logging.ERROR:
              log_level = trt.ILogger.Severity.ERROR
          elif level == logging.CRITICAL:
              log_level = trt.ILogger.Severity.INTERNAL_ERROR
          else:
              raise AssertionError(f"{level} is not valid log level")
  
          torch.ops.tensorrt.set_logging_level(int(log_level))
  
  
  def prepare_inputs(
      inputs: Input | torch.Tensor | Sequence[Any] | Dict[Any, Any],
      disable_memory_format_check: bool = False,
  ) -> Any:
      """
      We take a nested group of torch.Tensors or scalars and convert them into torchtrt.Input's
      """
      # Any tensors created inside this call will be FakeTensors if it's inside a torch.compile session
      # So, we disable fake mode temporarily.
      with unset_fake_temporarily():
          if inputs is None:
              return None
  
          elif isinstance(inputs, Input):
              return inputs
  
          elif isinstance(inputs, (torch.Tensor, int, float, bool)):
              return Input.from_tensor(
                  torch.tensor(inputs),
                  disable_memory_format_check=disable_memory_format_check,
              )
  
          elif isinstance(inputs, (list, tuple)):
              torchtrt_input_list = []
              for input_obj in inputs:
                  torchtrt_input = prepare_inputs(
                      input_obj, disable_memory_format_check=disable_memory_format_check
                  )
                  torchtrt_input_list.append(torchtrt_input)
  
              return (
                  torchtrt_input_list
                  if isinstance(inputs, list)
                  else tuple(torchtrt_input_list)
              )
  
          elif isinstance(inputs, dict):
              torchtrt_inputs_dict: Dict[Any, Any] = dict()
  
              for key, input_obj in inputs.items():
                  torchtrt_input = prepare_inputs(
                      input_obj, disable_memory_format_check=disable_memory_format_check
                  )
                  torchtrt_inputs_dict[key] = torchtrt_input
  
              return torchtrt_inputs_dict
  
          else:
              raise ValueError(
                  f"Invalid input type {type(inputs)} encountered in the dynamo_compile input parsing. "
                  + "Allowed input types: {torch_tensorrt.Input, torch.Tensor, list, tuple, dict}"
              )
  
  
  def parse_complex_tensor_structs(
      inputs: Input | torch.Tensor | Sequence[Any] | Dict[Any, Any],
      attribute_to_extract: str,
      apply_fn: Callable[[Any], Any] = lambda x: x,
  ) -> Any:
      """Parses complex structures of Tensors and returns a mirrored structure
      Extracts key attributes of each singular element, while reconstructing the struct
      Optionally applies a function to each attribute before returning
      """
      if isinstance(inputs, (torch.Tensor, Input)):
          return apply_fn(getattr(inputs, attribute_to_extract, None))
      elif isinstance(inputs, (int, float, bool)):
          # inputs is a python scalar value
          inputs_torch = torch.tensor(inputs)
          return apply_fn(getattr(inputs_torch, attribute_to_extract, None))
  
      elif isinstance(inputs, (list, tuple)):
          torchtrt_input_list = []
          for input_obj in inputs:
              torchtrt_input = parse_complex_tensor_structs(
                  input_obj, attribute_to_extract, apply_fn
              )
              torchtrt_input_list.append(torchtrt_input)
  
          return (
              torchtrt_input_list
              if isinstance(inputs, list)
              else tuple(torchtrt_input_list)
          )
  
      elif isinstance(inputs, dict):
          torchtrt_inputs_dict: Dict[Any, Any] = dict()
  
          for key, input_obj in inputs.items():
              torchtrt_input = parse_complex_tensor_structs(
                  input_obj, attribute_to_extract, apply_fn
              )
              torchtrt_inputs_dict[key] = torchtrt_input
  
          return torchtrt_inputs_dict
  
      else:
          raise ValueError(
              f"Invalid input type {type(inputs)} encountered during Dynamo input parsing. "
              + "Allowed input types: {torch_tensorrt.Input, torch.Tensor, list, tuple, dict}"
          )
  
  
  def contains_sym_int(tensor: torch.Tensor) -> bool:
      """
      Returns true if the given tensor has symbolic shape.
      """
      return any(isinstance(dim, torch.SymInt) for dim in tensor)
  
  
  def extract_var_range_info(symbolic_integer: torch.SymInt) -> Dict[str, int]:
      """
      This function returns the min, max, opt values of a symbolic integer.
      """
      node = symbolic_integer.node
      expr = node.expr
      shape_env = node.shape_env
      # An expr can be a independent SymInt node (eg: s0 or s1) or a composition of them eg: (48*s0 or s0*s1).
      # In the case of expr which has symbolic computation, bound_sympy evaluates them.
      # https://pytorch.org/docs/stable/generated/torch.fx.experimental.symbolic_shapes.ShapeEnv.html#torch.fx.experimental.symbolic_shapes.ShapeEnv.bound_sympy
      # expr.xreplace replaces the symbolic variables with their current values and computes the expression.
      var_range = shape_env.var_to_range.get(expr, None) or shape_env.bound_sympy(expr)
      var_val = (
          shape_env.var_to_val.get(expr, None)
          or shape_env.unbacked_var_to_val.get(expr, None)
          or expr.xreplace(shape_env.var_to_val)
      )
      assert var_range, var_val
      min_val, max_val = int(var_range.lower), int(var_range.upper)
      # Torchdynamo 0/1 specialization outlier
      min_val = 1 if min_val == 2 else min_val
      min_max_opt = {}
      min_max_opt["min"] = min_val
      min_max_opt["max"] = max_val
      if isinstance(var_val, (sympy.core.numbers.Integer, int)):
          min_max_opt["opt"] = int(var_val)
  
      return min_max_opt
  
  
  def unwrap_tensor_shape(
      tensor: Union[torch.Tensor, FakeTensor, torch.SymInt], mode: Optional[str] = ""
  ) -> Sequence[Union[int, Tuple[int, int]]]:
      """
      This is a helper function used to print/return the shape of the tensor.
      For regular torch.tensor's, it returns the static shape.
      For symbolic tensors, eg:(1, s0, 4), this function returns [1, [min, max], 4]. The min
      and max correspond to the lower and upper values of s0 symbolic dimension.
      """
      tensor_shape: List[Union[int, Tuple[int, int]]] = []
      # for dimension in tensor.shape:
      if isinstance(tensor, int):
          tensor_shape.append(tensor)
      elif isinstance(tensor, torch.SymInt):
          min_max_opt = extract_var_range_info(tensor)
          if mode:
              tensor_shape.append(min_max_opt[mode])
          else:
              tensor_shape.append((min_max_opt["min"], min_max_opt["max"]))
      elif isinstance(tensor, torch.SymFloat):
          # SymFloats can be an input to graph sometimes. Although SymFloat is scalar value, we treat it as a 1D tensor throughout Torch-TRT codebase.
          tensor_shape.append(1)
      elif isinstance(tensor, (torch.Tensor, FakeTensor)):
          for dimension in tensor.shape:
              tensor_shape.extend(unwrap_tensor_shape(dimension, mode=mode))
  
      return tuple(tensor_shape)
  
  
  def unwrap_tensor_dtype(tensor: Union[torch.Tensor, FakeTensor, torch.SymInt]) -> Any:
      """
      Returns the dtype of torch.tensor or FakeTensor. For symbolic integers, we return int64
      """
      if isinstance(tensor, (torch.Tensor, FakeTensor)):
          return tensor.dtype
      elif isinstance(tensor, (int, float, bool)):
          return torch.tensor(tensor).dtype
      elif isinstance(tensor, torch.SymInt):
          return torch.int64
      elif isinstance(tensor, torch.SymFloat):
          return torch.float32
      elif tensor is None:
          # Case where we explicitly pass one of the inputs to be None (eg: FLUX.1-dev)
          return None
      else:
          raise ValueError(f"Found invalid tensor type {type(tensor)}")
  
  
  def get_graph_io_attrs(
      io_nodes: Sequence[torch.fx.Node], attr_type: str
  ) -> Sequence[Any]:
      """
      Returns a list of attributes (shapes or dtypes) of the I/O nodes
      """
      assert attr_type in ["shape", "dtype"]
      attr_fn = unwrap_tensor_shape if attr_type == "shape" else unwrap_tensor_dtype
      graph_io_attrs = []
      for node in io_nodes:
          if "val" in node.meta:
              metadata = node.meta["val"]
              if isinstance(metadata, (tuple, list)):
                  for tensor in metadata:
                      graph_io_attrs.append(attr_fn(tensor))
              else:
                  graph_io_attrs.append(attr_fn(metadata))
  
      return graph_io_attrs
  
  
  def parse_graph_io(module: torch.fx.GraphModule, dryrun_tracker: Any) -> None:
      """
      Parse the graph I/O shape/dtype info for the whole graph and store in the dryrun tracker
      """
      # Parse inputs of the graph
      input_nodes = [node for node in module.graph.nodes if node.op == "placeholder"]
      input_shapes = get_graph_io_attrs(input_nodes, "shape")
      input_dtypes = get_graph_io_attrs(input_nodes, "dtype")
      dryrun_tracker.input_shapes = input_shapes
      dryrun_tracker.input_dtypes = input_dtypes
  
      # Parse outputs of the graph
      mark_output_nodes = [node for node in module.graph.nodes if node.op == "output"]
      output_nodes = []
      for node in mark_output_nodes:
          output_nodes.extend(node.all_input_nodes)
      output_shapes = get_graph_io_attrs(output_nodes, "shape")
      output_dtypes = get_graph_io_attrs(output_nodes, "dtype")
      dryrun_tracker.output_shapes = output_shapes
      dryrun_tracker.output_dtypes = output_dtypes
  
  
  def to_torch_device(device: Optional[Union[Device, torch.device, str]]) -> torch.device:
      """Cast a device-type to torch.device
  
      Returns the corresponding torch.device
      """
      if isinstance(device, Device):
          return device.to(torch.device)
  
      elif isinstance(device, torch.device):
          return device
  
      elif device is None:
          return torch.device(torch.cuda.current_device())
  
      else:
          return torch.device(device)
  
  
  def to_torch_tensorrt_device(
      device: Optional[Union[Device, torch.device, str]],
  ) -> Device:
      """Cast a device-type to torch_tensorrt.Device
  
      Returns the corresponding torch_tensorrt.Device
      """
      return Device._from(device)
  
  
  def parse_dynamo_kwargs(
      kwargs: Any,
  ) -> Tuple[CompilationSettings, Optional[BaseEngineCache]]:
      """Parses the kwargs field of a Dynamo backend
  
      Args:
          kwargs: Keyword arguments dictionary provided to the backend
      Returns:
          CompilationSettings object with relevant kwargs
      """
      # Initialize an empty CompilationSettings object
      settings = CompilationSettings()
  
      # If the user specifies keyword args, overwrite those fields in settings
      # Validate all specified kwargs to ensure they are true fields of the dataclass
      #
      # Note: kwargs provided by torch.compile are wrapped in the "options" key
      if kwargs:
          if "options" in kwargs and len(kwargs) == 1:
              kwargs = kwargs["options"]
  
          if "truncate_long_and_double" in kwargs:
              if (
                  "truncate_double" in kwargs
                  and kwargs["truncate_double"] is not _defaults.TRUNCATE_DOUBLE
              ):
                  raise ValueError(
                      'Provided configuration for "truncate_double" and deprecated API "truncate_long_and_double". '
                      'Please only use "truncate_double".'
                  )
              else:
                  kwargs["truncate_double"] = kwargs["truncate_long_and_double"]
                  warnings.warn(
                      'Compiler option "truncate_long_and_double" is deprecated in favor of "truncate_double" as int64 is now natively supported. '
                      "This option will be removed in the next version.",
                      DeprecationWarning,
                      stacklevel=2,
                  )
                  del kwargs[
                      "truncate_long_and_double"
                  ]  # Remove deprecated key after handling
  
          valid_attrs = {attr.name for attr in fields(settings)}
          valid_kwargs = {k: v for k, v in kwargs.items() if k in valid_attrs}
          settings = replace(settings, **valid_kwargs)
  
      # TODO: Remove once Dynamo precisions refactoring is complete
      if "enabled_precisions" in kwargs:
          enabled_precisions = {dtype._from(e) for e in kwargs["enabled_precisions"]}
  
          if len(enabled_precisions) == 0:
              logger.info(
                  f"No precision specified, defaulting to {_defaults.ENABLED_PRECISION}"
              )
              enabled_precisions = _defaults.ENABLED_PRECISIONS
  
          settings.enabled_precisions = enabled_precisions
  
      # Parse input runtime specification
      settings.use_python_runtime = use_python_runtime_parser(settings.use_python_runtime)
  
      # Ensure device is a torch_tensorrt Device
      settings.device = to_torch_tensorrt_device(settings.device)
  
      # Check and update device settings
      if "device" not in kwargs:
          logger.info(
              f"Device not specified, using Torch default current device - cuda:{settings.device.gpu_id}. "
              "If this is incorrect, please specify an input device, via the device keyword."
          )
  
      # Ignore and warn about require_full_compilation flag
      if settings.require_full_compilation:
          logger.warning(
              "Detected require_full_compilation=True for a torch.compile run. "
              "This option has no effect in torch.compile."
          )
          settings.require_full_compilation = False
  
      # If cache_built_engines and reuse_cached_engines are True but custom_engine_cache is not provided,
      # then create a default disk engine cache
  
      engine_cache = None
      if kwargs.get("cache_built_engines") or kwargs.get("reuse_cached_engines"):
          if kwargs.get("custom_engine_cache") is not None:
              engine_cache = kwargs.get("custom_engine_cache")
          else:
              from torch_tensorrt.dynamo._engine_cache import DiskEngineCache
  
              engine_cache_dir = kwargs.get(
                  "engine_cache_dir", _defaults.ENGINE_CACHE_DIR
              )
              engine_cache_size = kwargs.get(
                  "engine_cache_size", _defaults.ENGINE_CACHE_SIZE
              )
              engine_cache = DiskEngineCache(engine_cache_dir, engine_cache_size)
  
      if kwargs.get("torch_executed_ops"):
          settings.torch_executed_ops = kwargs.get("torch_executed_ops")
  
      logger.info("Compilation Settings: %s\n", settings)
  
      return settings, engine_cache
  
  
  def req_torch_version(min_torch_version: str = "2.dev") -> Callable[..., Any]:
      """
      Create a decorator which verifies the Torch version installed
      against a specified version range
  
      Args:
          min_torch_version (str): The minimum required Torch version
          for the decorated function to work properly
  
      Returns:
          A decorator which raises a descriptive error message if
          an unsupported Torch version is used
      """
  
      def nested_decorator(f: Callable[..., Any]) -> Callable[..., Any]:
          def function_wrapper(*args: Any, **kwargs: Any) -> Any:
              # Parse minimum and current Torch versions
              min_version = version.parse(min_torch_version)
              current_version = version.parse(torch.__version__)
  
              if current_version < min_version:
                  raise AssertionError(
                      f"Expected Torch version {min_torch_version} or greater, "
                      + f"when calling {f}. Detected version {torch.__version__}"
                  )
              else:
                  return f(*args, **kwargs)
  
          return function_wrapper
  
      return nested_decorator
  
  
  def check_module_output(
      new_module: torch.fx.GraphModule,
      refitted_module: torch.fx.GraphModule,
      arg_inputs: Any,
      kwarg_inputs: Any = None,
  ) -> bool:
      old_outputs, new_outputs = refitted_module(*arg_inputs), new_module(
          *arg_inputs, **kwarg_inputs
      )
      if type(old_outputs) != type(new_outputs):
          logger.warning("The output types are different. Output check is skipped.")
          return True
      return check_output_equal(old_outputs, new_outputs)
  
  
  def check_output_equal(
      output1: Any,
      output2: Any,
      rtol: float = RTOL,
      atol: float = ATOL,
  ) -> bool:
      if type(output1) != type(output2):
          logger.warning(
              "The output types are different. Check_output_equal will always return false."
          )
          return False
  
      if isinstance(output1, torch.Tensor):
          if output1.shape != output2.shape:
              return False
          return torch.allclose(output1, output2, rtol, atol)  # type: ignore
  
      elif isinstance(output1, (tuple, list)):
          if len(output1) != len(output2):
              return False
          for a, b in zip(output1, output2):
              if not check_output_equal(a, b):
                  return False
              return True
  
      elif isinstance(output1, dict):
          if output1.keys() != output2.keys():
              return False
          for a, b in zip(output1.values(), output2.values()):
              if not check_output_equal(a, b):
                  return False
          return True
  
      logger.warning(
          "The output type is not supported to be checked. Check_output_equal will always return false."
      )
      return False
  
  
  def get_flat_args_with_check(
      exported_program: torch.export.ExportedProgram,
      args: list[Any],
      kwargs: dict[str, Any],
  ) -> tuple[Any, Any]:
      """Flatten args, kwargs using pytree, then, check specs.
  
      Args:
          args: List[Any] original args passed to __call__
          kwargs: Dict[str, Any] original kwargs passed to __call
  
      Returns:
          A tuple of (flat_args, received_spec)
          flat_args is flattend args / kwargs
          received_spec is the pytree spec produced while flattening the
          tuple (args, kwargs)
      """
      import torch.utils._pytree as pytree
      from torch.export._tree_utils import reorder_kwargs
  
      in_spec = exported_program.call_spec.in_spec
      if in_spec is not None:
          kwargs = reorder_kwargs(kwargs, in_spec)
      flat_args_with_path, received_spec = pytree.tree_flatten_with_path((args, kwargs))
      flat_args = tuple(x[1] for x in flat_args_with_path)
      return flat_args, received_spec
  
  
  def get_metadata(
      gm: torch.fx.GraphModule, target_op: torch._ops.OpOverload
  ) -> List[Any]:
      """
      Return the list which has the metadata of all the target_op nodes present in the graph.
      """
      return [node.meta for node in gm.graph.nodes if node.target == target_op]
  
  
  def set_metadata(
      gm: torch.fx.GraphModule, target_op: torch._ops.OpOverload, metadata: List[Any]
  ) -> None:
      """
      Return the list which has the metadata of all the target_op nodes present in the graph.
      """
      target_nodes = [node for node in gm.graph.nodes if node.target == target_op]
      assert len(target_nodes) == len(metadata)
      for idx, node in enumerate(target_nodes):
          node.meta = metadata[idx]
  
  
  def copy_metadata(match_and_replacements: List[Any]) -> None:
      """
      Copy the metadata from anchor node to the replacement node. This should be used
      if the anchor node is replaced with only a single replacement node i.e one-one replacement.
      """
      for match_and_replacement in match_and_replacements:
          anchor_node = match_and_replacement.nodes_map[match_and_replacement.anchor]
          assert (
              len(match_and_replacement.replacements) == 1
          ), "Found more than 1 replacements for the anchor node."
          replacement_node = match_and_replacement.replacements[0]
          replacement_node.meta = anchor_node.meta
  
  
  def flatten_nodes(nodes: Any) -> List[torch.fx.node.Node]:
      ret = []
      if isinstance(nodes, torch.fx.node.Node):
          ret.append(nodes)
      elif isinstance(nodes, (tuple, list)):
          for node in nodes:
              ret.extend(flatten_nodes(node))
      else:
          raise ValueError(
              f"expect torch.fx.node.Node or a tuple/list of torch.fx.node.Node type, got unexpected types: {type(nodes)=}"
          )
      return ret
  
  
  def get_output_metadata(
      gm: torch.fx.GraphModule,
  ) -> List[Any]:
      outputs = [node for node in gm.graph.nodes if node.op == "output"]
      assert len(outputs) > 0
      outputs = outputs[0].args
      nodes = flatten_nodes(outputs)
      assert len(nodes) > 0
      return [node.meta for node in nodes]
  
  
  def get_output_dtypes(output: Any, truncate_double: bool = False) -> List[dtype]:
      output_dtypes = []
      if isinstance(output, torch.fx.node.Node):
          if "val" in output.meta:
              output_meta = output.meta["val"]
              if isinstance(output_meta, (FakeTensor, torch.Tensor)):
                  if truncate_double and output_meta.dtype == torch.float64:
                      output_dtypes.append(dtype.float32)
                  else:
                      output_dtypes.append(dtype._from(output_meta.dtype))
              elif isinstance(output_meta, torch.SymInt):
                  output_dtypes.append(dtype.int64)
          elif "tensor_meta" in output.meta:
              output_meta = output.meta["tensor_meta"]
              output_dtypes.append(dtype._from(output_meta.dtype))
          else:
              raise ValueError(
                  f"node.name={output.name}: metadata does not exist, expect metadata exists for each output node"
              )
      elif isinstance(output, (tuple, list)):
          for ele in output:
              output_dtypes.extend(get_output_dtypes(ele))
      else:
          raise ValueError(
              f"got unexpected type {type(output)}, expected type is a torch.fx.node.Node or a tuple/list of torch.fx.node.Node"
          )
      return output_dtypes
  
  
  def get_cpu_memory_usage() -> Any:
      return psutil.Process().memory_info().rss / 1024 / 1024
  
  
  def release_host_and_device_memory() -> None:
      gc.collect()
      if torch.cuda.is_available():
          torch.cuda.synchronize()
          torch.cuda.empty_cache()
          torch.cuda.ipc_collect()
          torch.cuda.synchronize()
  
      if (
          platform.system() == "Linux"
          and os.environ.get("TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM", "0") == "1"
      ):
          try:
              libc = ctypes.CDLL("libc.so.6")
              if libc.malloc_trim(0) != 1:
                  logger.warning("Failed to release CPU memory.")
          except Exception:
              logger.warning("Failed to release CPU memory.")
  