content: from __future__ import annotations
  
  import collections.abc
  import copy
  import gc
  import logging
  from typing import Any, List, Optional, Sequence, Tuple
  
  import numpy as np
  import tensorrt as trt
  import torch
  from torch.export import ExportedProgram
  from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
  from torch_tensorrt._enums import dtype
  from torch_tensorrt._features import needs_refit
  from torch_tensorrt._Input import Input
  from torch_tensorrt.dynamo import partitioning
  from torch_tensorrt.dynamo._exporter import inline_torch_modules
  from torch_tensorrt.dynamo._settings import CompilationSettings
  from torch_tensorrt.dynamo.conversion._conversion import infer_module_output_dtypes
  from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
      DYNAMO_CONVERTERS as CONVERTERS,
  )
  from torch_tensorrt.dynamo.conversion._TRTInterpreter import TRTInterpreter
  from torch_tensorrt.dynamo.conversion.impl.normalization.ops import (
      batch_norm_constant_folding,
  )
  from torch_tensorrt.dynamo.conversion.truncate_double import repair_double_inputs
  from torch_tensorrt.dynamo.lowering import (
      clean_up_graph_after_modifications,
      get_decompositions,
      post_lowering,
      pre_export_lowering,
  )
  from torch_tensorrt.dynamo.runtime._PythonTorchTensorRTModule import (
      PythonTorchTensorRTModule,
  )
  from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (
      ENGINE_IDX,
      SERIALIZED_METADATA_IDX,
      TorchTensorRTModule,
  )
  from torch_tensorrt.dynamo.utils import (
      CPU_DEVICE,
      check_module_output,
      get_model_device,
      get_torch_inputs,
      to_torch_device,
      to_torch_tensorrt_device,
  )
  from torch_tensorrt.logging import TRT_LOGGER
  
  logger = logging.getLogger(__name__)
  
  
  @needs_refit  # type: ignore[misc]
  def construct_refit_mapping(
      module: torch.fx.GraphModule,
      inputs: Sequence[Input],
      settings: CompilationSettings = CompilationSettings(),
  ) -> dict[str, np.ndarray]:
      """Find out the weight mapping between weight in exported program and TensorRT engine
      Args:
          module: FX GraphModule to interpret
          inputs: Sequence of Tensors representing inputs to the module
          settings: Compilation settings
      Returns:
          Mapping from weight name in TensorRT to actual weight value in np.ndarray
      """
  
      output_dtypes = infer_module_output_dtypes(
          module,
          truncate_double=settings.truncate_double,
      )
  
      # Use Interpreter
      interpreter = TRTInterpreter(
          module,
          inputs,
          output_dtypes=output_dtypes,
          compilation_settings=settings,
      )
      interpreter._construct_trt_network_def()
      weight_refit_map: dict[str, torch.Tensor] = interpreter.ctx.weight_refit_map
  
      return weight_refit_map
  
  
  @needs_refit  # type: ignore[misc]
  def construct_refit_mapping_from_weight_name_map(
      weight_name_map: dict[Any, Any],
      state_dict: dict[Any, Any],
      settings: CompilationSettings,
  ) -> dict[Any, Any]:
      engine_weight_map = {}
      for engine_weight_name, (sd_weight_name, np_weight_type) in weight_name_map.items():
          # Add more constant folding converters here
          trt_dtype = dtype._from(np_weight_type).to(trt.DataType)
          torch_dtype = dtype._from(np_weight_type).to(torch.dtype)
          if engine_weight_name.split(" ")[-1] in ["SCALE", "SHIFT"]:
              # Batch Norm Layer
              params = {}
              for w in sd_weight_name:
                  params[w.split(".")[-1]] = state_dict[w].cuda()
              # Batch norm constant folding
  
              scale, shift = batch_norm_constant_folding(**params, eps=1e-5)
              # Set scale to scale or shift to shift
              engine_weight_map[engine_weight_name] = eval(
                  engine_weight_name.split(" ")[-1].lower()
              )
  
          elif sd_weight_name not in state_dict:
              # If weights is not in sd, we can leave it unchanged
              continue
          else:
  
              engine_weight_map[engine_weight_name] = state_dict[sd_weight_name].to(
                  to_torch_device(settings.device)
              )
  
          engine_weight_map[engine_weight_name] = (
              engine_weight_map[engine_weight_name]
              .clone()
              .reshape(-1)
              .contiguous()
              .to(torch_dtype),
              trt_dtype,
          )
  
      return engine_weight_map
  
  
  @needs_refit  # type: ignore[misc]
  def _refit_single_trt_engine_with_gm(
      new_gm: torch.fx.GraphModule,
      old_engine: trt.ICudaEngine,
      input_list: Sequence[Any],
      settings: CompilationSettings = CompilationSettings(),
      weight_name_map: Optional[dict[str, List[str]]] = None,
  ) -> None:
      """
      Refit a TensorRT Engine in place
      """
  
      with unset_fake_temporarily():
          refitted = set()
          torch_device = get_model_device(new_gm)
          refitter = trt.Refitter(old_engine, TRT_LOGGER)
          weight_list = refitter.get_all_weights()
  
          if weight_name_map:
              # Get the refitting mapping
              trt_wt_location = (
                  trt.TensorLocation.DEVICE
                  if torch_device.type == "cuda"
                  else trt.TensorLocation.HOST
              )
  
              constant_mapping: dict[str, Any] = weight_name_map.pop(
                  "constant_mapping", {}
              )  # type: ignore
              mapping = construct_refit_mapping_from_weight_name_map(
                  weight_name_map, new_gm.state_dict(), settings
              )
              constant_mapping_with_type = {}
  
              for constant_name, val in constant_mapping.items():
                  np_weight_type = val.dtype
                  val_tensor = torch.from_numpy(val).cuda()
                  trt_dtype = dtype._from(np_weight_type).to(trt.DataType)
                  torch_dtype = dtype._from(np_weight_type).to(torch.dtype)
                  constant_mapping_with_type[constant_name] = (
                      val_tensor.clone().reshape(-1).contiguous().to(torch_dtype),
                      trt_dtype,
                  )
  
              mapping.update(constant_mapping_with_type)
  
              for layer_name in weight_list:
                  if layer_name not in mapping:
                      logger.warning(f"{layer_name} is not found in weight mapping.")
                      continue
                  # Use Numpy to create weights
                  weight, weight_dtype = mapping[layer_name]
                  trt_wt_tensor = trt.Weights(
                      weight_dtype, weight.data_ptr(), torch.numel(weight)
                  )
                  refitter.set_named_weights(layer_name, trt_wt_tensor, trt_wt_location)
              assert (
                  len(refitter.get_missing_weights()) == 0
              ), "Fast refitting failed due to incomplete mapping"
  
          else:
              mapping = construct_refit_mapping(new_gm, input_list, settings)
              trt_wt_location = trt.TensorLocation.HOST
              for layer_name in weight_list:
                  if layer_name not in mapping:
                      raise AssertionError(f"{layer_name} is not found in weight mapping")
                  # Use Tensor to create weights
                  weight = mapping[layer_name]
                  trt_dtype = dtype._from(weight.dtype).to(trt.DataType)
                  trt_wt_tensor = trt.Weights(
                      trt_dtype, weight.data_ptr(), torch.numel(weight)
                  )
                  refitter.set_named_weights(layer_name, trt_wt_tensor, trt_wt_location)
                  refitted.add(layer_name)
  
              if len(refitted) != len(weight_list):
                  logger.warning("Not all weights have been refitted!!!")
  
          if not refitter.refit_cuda_engine():
              logger.error("Error: failed to refit new weights.")
              raise AssertionError("Refitting failed.")
  
  
  @needs_refit  # type: ignore[misc]
  def refit_module_weights(
      compiled_module: torch.fx.GraphModule | ExportedProgram,
      new_weight_module: ExportedProgram,
      arg_inputs: Optional[Tuple[Any, ...]] = None,
      kwarg_inputs: Optional[dict[str, Any]] = None,
      verify_output: bool = False,
      use_weight_map_cache: bool = True,
      in_place: bool = False,
  ) -> torch.fx.GraphModule:
      """
      Refit a compiled graph module with ExportedProgram. This performs weight updates in compiled_module without recompiling the engine.
  
      Args:
          compiled_module: compiled TensorRT module that needs to be refitted.
                          This compiled_module should be compmiled by torch_tensorrt.dynamo.compile
                          or load it from disk using trt.load.
          new_weight_module: exported program with the updated weights. This one should have the same model architecture as the compiled module.
          arg_inputs: sample arg inputs. Optional, needed if output check
          kwarg_inputs: sample kwarg inputs. Optional, needed if output check
          verify_output: whether to verify output of refitted module
      Returns:
          A new compiled TensorRT module that has the updated weights.
      """
      inline_module = False
      if isinstance(compiled_module, ExportedProgram):
          compiled_module = compiled_module.module()
  
      if len(list(compiled_module.named_children())) == 0:
          inline_module = True
  
      if not in_place:
          compiled_module = copy.deepcopy(compiled_module)
      elif inline_module:
          raise AssertionError(
              "Exported program does not support modifying in place. Please set in_place to false and use the returned graph module."
          )
  
      # Get the settings and check the setting to be uniform
      settings: Optional[CompilationSettings] = None
      if inline_module:
          # Obtain the settings
          compiled_submodules = [
              (name.replace("_engine", ""), engine)
              for name, engine in compiled_module.__dict__.items()
              if "engine" in name
          ]
          # [('_run_on_acc_0', inline_module)]
          encoded_metadata = compiled_submodules[0][1].__getstate__()[0][
              SERIALIZED_METADATA_IDX
          ]
          assert (
              encoded_metadata != ""
          ), "The engine provided is either not refittable or was built with a version of Torch-TensorRT that is too old, please recompile using the latest version"
          settings = TorchTensorRTModule.decode_metadata(encoded_metadata)["settings"]
          # Handle torch modules
          compiled_submodules_map = dict(compiled_submodules)
          for name, submodule in compiled_module.named_children():
              compiled_submodules_map[name] = submodule
  
      else:
          # Handle torch modules
          compiled_submodules_map = {}
          guard_fn_modules = []
          for name, submodule in compiled_module.named_children():
              if (
                  not isinstance(
                      submodule,
                      (
                          PythonTorchTensorRTModule,
                          TorchTensorRTModule,
                          torch.nn.modules.module.Module,
                      ),
                  )
                  or "_run_on_gpu" in name
              ):
                  continue
  
              # When we re-export the graph module, torch.export._unlift.GuardsFn modules are being added as submodules.
              if isinstance(submodule, torch.export._unlift.GuardsFn):
                  guard_fn_modules.append(name)
                  continue
              # Obtain the settings
  
              compiled_submodules = [
                  (name.replace("_engine", ""), engine)
                  for name, engine in submodule.__dict__.items()
                  if "engine" in name
              ]
  
              settings = None
              try:
                  # If the gm is not inlined or transformed by retracing, the settings is stored in the submodule
                  settings = submodule.settings
              except AttributeError:
  
                  encoded_metadata = [
                      engine for name, engine in compiled_submodules if name == "engine"
                  ][0].__getstate__()[0][SERIALIZED_METADATA_IDX]
                  assert (
                      encoded_metadata != ""
                  ), "The engine provided is either not refittable or was built with a version of Torch-TensorRT that is too old, please recompile using the latest version"
                  settings = TorchTensorRTModule.decode_metadata(encoded_metadata)[
                      "settings"
                  ]
  
              compiled_submodules_map[name] = submodule
  
          # Delete the guard fn modules to avoid the guard fn modules being refitted
          # First, remove nodes in the graph that reference the guard function modules
          for node in list(compiled_module.graph.nodes):
              if node.op == "call_module" and node.target in guard_fn_modules:
                  compiled_module.graph.erase_node(node)
  
          # Now delete the submodules themselves
          for guard_fn_module_name in guard_fn_modules:
              # delattr(compiled_module, guard_fn_module_name)
              compiled_module.delete_submodule(guard_fn_module_name)
  
          # Clean up the graph
          clean_up_graph_after_modifications(compiled_module)
  
      assert settings is not None
  
      assert (
          not settings.immutable_weights
      ), "Refitting is not enabled. Please recompile the engine with immutable_weights=False."
  
      device = to_torch_tensorrt_device(settings.device)
      if arg_inputs:
          if not isinstance(arg_inputs, collections.abc.Sequence):
              # Prepare torch_trt inputs
              arg_inputs = [arg_inputs]
          torch_inputs = get_torch_inputs(arg_inputs, device)
  
      torch_kwarg_inputs: Any = {}
      if kwarg_inputs:
          torch_kwarg_inputs = get_torch_inputs(kwarg_inputs, device)
      runtime = trt.Runtime(TRT_LOGGER)
      if not isinstance(new_weight_module, ExportedProgram):
          raise AssertionError(
              f"Input graph should be an ExportedProgram but got type {type(new_weight_module)}"
          )
      new_weight_module = pre_export_lowering(new_weight_module, settings)
      new_weight_module = new_weight_module.run_decompositions(
          get_decompositions(settings.enable_experimental_decompositions)
      )
      new_gm = new_weight_module.module()
  
      logger.debug("Input graph: " + str(new_gm.graph))
      # Apply lowering on the graph module
  
      new_gm = post_lowering(new_gm, settings)
  
      logger.debug("Lowered Input graph: " + str(new_gm.graph))
  
      # Set torch-executed ops
      CONVERTERS.set_compilation_settings(settings)
  
      # Check the number of supported operations in the graph
      num_supported_ops, total_ops = partitioning.get_graph_converter_support(
          new_gm, settings.torch_executed_ops
      )
  
      if num_supported_ops == 0 or (
          num_supported_ops < settings.min_block_size and not settings.dryrun
      ):
          logger.warning(
              f"{num_supported_ops} supported operations detected in subgraph containing {total_ops} computational nodes. "
              f"Skipping this subgraph, since min_block_size was detected to be {settings.min_block_size}"
          )
          return new_gm
      else:
          logger.debug(
              f"Detected support for {num_supported_ops} operators out of {total_ops} in subgraph."
          )
  
      # If specified, try using the fast partitioner and fall back to the global one on failure
      if settings.use_fast_partitioner:
          try:
              logger.info("Partitioning the graph via the fast partitioner")
              new_partitioned_module, supported_ops = partitioning.fast_partition(
                  new_gm,
                  min_block_size=settings.min_block_size,
                  torch_executed_ops=settings.torch_executed_ops,
                  require_full_compilation=settings.require_full_compilation,
                  skip_fusion=(num_supported_ops == total_ops),
              )
  
          except torch.fx.passes.splitter_base.FxNetSplitterInternalError:
              logger.error(
                  "Partitioning failed on the subgraph with fast partition. See trace above. "
                  "Retrying with global partition.",
                  exc_info=True,
              )
  
              settings.use_fast_partitioner = False
  
      if not settings.use_fast_partitioner:
          logger.info("Partitioning the graph via the global partitioner")
          new_partitioned_module, supported_ops = partitioning.global_partition(
              new_gm,
              min_block_size=settings.min_block_size,
              torch_executed_ops=settings.torch_executed_ops,
              require_full_compilation=settings.require_full_compilation,
          )
  
      # Done Partition
      if inline_module:
          # Preprocess the partitioned module to be in the same format as the inline module
          inline_torch_modules(new_partitioned_module)
          new_partitioned_module.delete_all_unused_submodules()
          # Check the number of partitions and name
          assert {sm[0] for sm in new_partitioned_module.named_children()} == set(
              compiled_submodules_map.keys()
          ), "New weights module is not compatible with previously compiled Torch-TensorRT module"
      else:
          assert {sm[0] for sm in new_partitioned_module.named_children()} == {
              sm[0] for sm in compiled_module.named_children()
          }, "New weights module is not compatible with previously compiled Torch-TensorRT module"
      # 2. TODO: Check the hash of source fx.Graph and new fx.Graph
  
      # Iterate over all components that can be accelerated
      # Generate the corresponding TRT Module for those
      for name, new_submodule in new_partitioned_module.named_children():
          # Refit each submodule
          # Extract engine from the submodule
          try:
              if inline_module:
                  weight_name_map = None
                  compiled_submodule = compiled_submodules_map[name]
                  # If this is a torch module, load the old state_dict
                  if "_run_on_acc" not in name:
                      compiled_submodule.load_state_dict(new_submodule.state_dict())
                      continue
                  else:
                      engine_info = compiled_submodule.__getstate__()[0]
                      engine = get_engine_from_encoded_engine(
                          engine_info[ENGINE_IDX], runtime
                      )
                      if use_weight_map_cache:
                          encoded_metadata = compiled_submodule.__getstate__()[0][
                              SERIALIZED_METADATA_IDX
                          ]
                          weight_name_map = TorchTensorRTModule.decode_metadata(
                              encoded_metadata
                          )["weight_name_map"]
                          if not weight_name_map:
                              use_weight_map_cache = False
                              logger.warning(
                                  "This engine does not have a weight map cache. Rebuilding the weight map"
                              )
              else:
                  compiled_submodule = getattr(compiled_module, name)
                  if "_run_on_acc" not in name:
                      compiled_submodule.load_state_dict(new_submodule.state_dict())
                      continue
  
                  weight_name_map = None
                  if use_weight_map_cache:
                      try:
                          weight_name_map = compiled_submodule.weight_name_map
                      except AttributeError:
                          if isinstance(compiled_submodule, torch.nn.Module):
                              # Torch retrace module
                              assert (
                                  not settings.use_python_runtime
                              ), "Refitting a torch retraced module is only supported with use_python_runtime=False"
                              encoded_metadata = [
                                  engine
                                  for name, engine in compiled_submodules
                                  if name == "engine"
                              ][0].__getstate__()[0][SERIALIZED_METADATA_IDX]
                              weight_name_map = TorchTensorRTModule.decode_metadata(
                                  encoded_metadata
                              )["weight_name_map"]
  
                          if not isinstance(
                              compiled_submodule, torch.fx.graph_module.GraphModule
                          ):
                              logger.warning(
                                  "The module was compiled with an old version of Torch-TensorRT. Rebuilding the weight map."
                              )
                      if not weight_name_map:
                          use_weight_map_cache = False
                          logger.warning(
                              "This engine does not have a weight map cache. Rebuilding the weight map"
                          )
  
                  # Rexporting the TRT compiled graph module and loading it back doesn't preserve the instance type and registers
                  # the compiled submodule as torch.nn.Module. So we use settings.use_python_runtime to determine the instance type.
                  if settings.use_python_runtime:
                      engine = compiled_submodule.engine
                  else:
                      engine_info = compiled_submodule.engine.__getstate__()[0]
                      engine = get_engine_from_encoded_engine(
                          engine_info[ENGINE_IDX], runtime
                      )
          except AttributeError:
              raise AssertionError(
                  "The type of graph module is not supported for refitting or two compiled modules do not match."
              )
  
          # Get the submodule inputs for min, opt, max shapes of the graph inputs
          submodule_inputs = partitioning.construct_submodule_inputs(new_submodule)
          logger.debug(
              "Refitting Submodule name: %s\n",
              str(name),
          )
          assert submodule_inputs is not None
          # Handle long/double inputs if requested by the user
          if settings.truncate_double:
              submodule_inputs = repair_double_inputs(
                  new_partitioned_module,
                  new_submodule,
                  submodule_inputs,
                  to_torch_device(settings.device),
                  name,
              )
          try:
              _refit_single_trt_engine_with_gm(
                  new_gm=new_submodule,
                  old_engine=engine,
                  input_list=submodule_inputs,
                  settings=settings,
                  weight_name_map=weight_name_map,
              )
  
          except AssertionError as e:
              # If fast_refit is used and failed, we fall back to regular refit
              logger.warning(e)
              if use_weight_map_cache and weight_name_map:
                  _refit_single_trt_engine_with_gm(
                      new_gm=new_submodule,
                      old_engine=engine,
                      input_list=submodule_inputs,
                      settings=settings,
                      weight_name_map=None,
                  )
  
          # clear EXCLUDE_WEIGHTS flag and set INCLUDE_REFIT flag to make the engine refittable
          serialization_config = engine.create_serialization_config()
          serialization_config.clear_flag(trt.SerializationFlag.EXCLUDE_WEIGHTS)
          # INCLUDE_REFIT flag is only available as of TensorRT 10.14
          if hasattr(trt.SerializationFlag, "INCLUDE_REFIT"):
              serialization_config.set_flag(trt.SerializationFlag.INCLUDE_REFIT)
          serialized_engine = engine.serialize_with_config(serialization_config)
  
          if isinstance(compiled_submodule, PythonTorchTensorRTModule):
              compiled_submodule.serialized_engine = bytes(serialized_engine)
          elif isinstance(compiled_submodule, TorchTensorRTModule):
              compiled_submodule.engine = None  # Clear the engine for TorchTensorRTModule, otherwise it won't be updated
              compiled_submodule.serialized_engine = bytes(serialized_engine)
              compiled_submodule.setup_engine()
          elif inline_module:
              new_engine_info = list(engine_info)
              new_engine_info[ENGINE_IDX] = bytes(serialized_engine)
              refitted_engine = torch.classes.tensorrt.Engine(tuple(new_engine_info))
              setattr(compiled_module, f"{name}_engine", refitted_engine)
          elif isinstance(compiled_submodule, torch.nn.Module):
              # Torch retrace module
              new_engine_info = list(engine_info)
              new_engine_info[ENGINE_IDX] = bytes(serialized_engine)
              refitted_engine = torch.classes.tensorrt.Engine(tuple(new_engine_info))
              compiled_submodule.engine = refitted_engine
          del engine
          gc.collect()
          torch.cuda.empty_cache()
  
      if verify_output and arg_inputs is not None:
          new_gm.to(to_torch_device(settings.device))
          if check_module_output(
              new_module=new_gm,
              refitted_module=compiled_module,
              arg_inputs=torch_inputs,
              kwarg_inputs=torch_kwarg_inputs,
          ):
              logger.info("Refitting Succeed!")
              new_gm.to(CPU_DEVICE)
          else:
              if weight_name_map:
                  logger.warning(
                      "Refitting with weight_name_map yielded incorrect result! The outputs do not match."
                  )
                  return refit_module_weights(
                      compiled_module,
                      new_weight_module,
                      arg_inputs,
                      kwarg_inputs,
                      verify_output,
                      use_weight_map_cache=False,
                      in_place=in_place,
                  )
              logger.error("Refitting Failed! The outputs do not match.")
              new_gm.to(CPU_DEVICE)
      else:
          logger.info("Refitting Completed! Output verification skipped.")
  
      return compiled_module
  
  
  # Util functions -----------
  import base64
  
  
  def get_engine_from_encoded_engine(
      encoded_engine: str, runtime: trt.Runtime
  ) -> trt.ICudaEngine:
      serialized_engine = base64.b64decode(encoded_engine)
      engine = runtime.deserialize_cuda_engine(serialized_engine)
      return engine
  