content: from __future__ import annotations
  
  import logging
  from typing import Any, List, Optional, Sequence, Tuple
  
  import torch
  import torch_tensorrt
  from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
  from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
  from torch_tensorrt.dynamo import partitioning
  
  logger = logging.getLogger(__name__)
  
  
  def _unflatten_inputs(
      flattened_inputs: Sequence[torch_tensorrt.Input],
      compiled_module: torch.fx.GraphModule,
  ) -> Tuple[Any, Any]:
      """
      Process inputs using tree_unflatten and tree_map to reconstructe inputs
  
      Args:
          flattened_inputs: Flattened input tensors to process
          compiled_module: The compiled GraphModule containing input specifications
  
      Returns:
          Tuple of (args, kwargs) containing reconstructed input tensors
      """
  
      def convert_input_to_cuda_tensor(input: Any) -> torch.Tensor:
          if isinstance(input, torch_tensorrt.Input):
              return input.torch_tensor.cuda()
          else:
              raise RuntimeError("Input is not a torch_tensorrt.Input")
  
      # Reconstruct the (args, kwargs) structure that was flattened during export
      pytree_inputs = tree_unflatten(flattened_inputs, compiled_module._in_spec)
      # Apply the tensor creation to the reconstructed structure
      processed_inputs = tree_map(convert_input_to_cuda_tensor, pytree_inputs)
  
      # Since inputs were originally flattened from (args, kwargs),
      # processed_inputs is now that same tuple structure
      return processed_inputs[0], processed_inputs[1]
  
  
  class CudaGraphsTorchTensorRTModule(torch.nn.Module):  # type: ignore[misc]
      """This Wrapper runtime module is to record/replay whole cuda graph in sub modules
  
      Args:
          compiled_module: Complied fx graphModule that will be wrapped
      Returns:
          Output tensor or tensor list
      """
  
      def __init__(
          self,
          compiled_module: torch.nn.Module,
      ):
          super(CudaGraphsTorchTensorRTModule, self).__init__()
          self.compiled_module = compiled_module
          self.inputs = partitioning.construct_submodule_inputs(compiled_module)
          self.is_weight_streaming_set = False
  
          self._input_buffers: List[torch.Tensor] = []
          self._output_buffers: List[torch.Tensor] = []
          self.cudagraph: Optional[torch.cuda.CUDAGraph] = None
          self.use_output_allocator_outputs = False
          self.shape_key: Optional[str] = None
          self._caller_stream: Optional[torch.cuda.Stream] = None
          self._engine_stream: Optional[torch.cuda.Stream] = None
          self.warm_up()
  
      def warm_up(self) -> None:
          """
          Warm up is necessary to ensure that memory allocations and initializations
          are not recorded in cuda graphs
          """
  
          with torch_tensorrt.logging.errors():
              with unset_fake_temporarily():
                  args, kwargs = _unflatten_inputs(self.inputs, self.compiled_module)
                  s = torch.cuda.Stream()
                  s.wait_stream(torch.cuda.current_stream())
                  with torch.cuda.stream(s):
                      for _ in range(3):
                          self.compiled_module(*args, **kwargs)
                  torch.cuda.current_stream().wait_stream(s)
  
      def validate_input_shapes(self, inputs: Sequence[torch.Tensor]) -> bool:
          """
          Validates the input shapes of the forward function has changed
          And infer output shapes if dynamic input shape has changed.
          """
          # Representation of input shapes to a given model
          # Shapes are concatenated as so:
          # x: (3, 4), y: (4, 5) --> Key: (3,4)(4,5)
          new_shape_key = "".join(str(tuple(t.shape)).replace(" ", "") for t in inputs)
  
          if new_shape_key != self.shape_key:
              logger.debug(f"Input shape changed {self.shape_key} -> {new_shape_key}")
              self.shape_key = new_shape_key
              return True
  
          return False
  
      def _reset_captured_graph(self) -> None:
          if self.cudagraph:
              self.cudagraph.reset()
              self.cudagraph = None
  
      def __del__(self) -> None:
          self._reset_captured_graph()
  
      def set_use_output_allocator(self, enable: bool) -> None:
          self.use_output_allocator_outputs = enable
  
      def forward(
          self, *args: Any, **kwargs: Any
      ) -> torch.Tensor | Tuple[torch.Tensor, ...]:
          inputs, _ = tree_flatten((args, kwargs))
          cudagraphs_enabled = torch_tensorrt.runtime.get_whole_cudagraphs_mode()
          if cudagraphs_enabled:
              shape_changed = self.validate_input_shapes(inputs)
              need_cudagraphs_record = shape_changed or self.is_weight_streaming_set
              if need_cudagraphs_record:
                  self._reset_captured_graph()
                  self._input_buffers = [None] * len(inputs)
  
              self.is_weight_streaming_set = False
              # Ensure inputs are available in all scopes and cast symbolic integers to Tensors
              contiguous_inputs: List[torch.Tensor] = [
                  (
                      i.contiguous()
                      if isinstance(i, torch.Tensor)
                      else torch.tensor(i).cuda()
                  )
                  for i in inputs
              ]
              assert len(contiguous_inputs) == len(
                  inputs
              ), f"Wrong number of inputs, expect {len(inputs)} get {len(contiguous_inputs)}."
  
              for i, _ in enumerate(inputs):
                  if not contiguous_inputs[i].is_cuda:
                      logger.warning(
                          f"Detected input[{i}] is not on a cuda device. "
                          "This tensor is being moved by the runtime but for performance considerations, "
                          "ensure your inputs are all on GPU and open an issue here "
                          "(https://github.com/pytorch/TensorRT/issues) if this warning persists."
                      )
                      contiguous_inputs = (
                          contiguous_inputs[:i]
                          + [contiguous_inputs[i].cuda()]
                          + contiguous_inputs[i + 1 :]
                      )
  
                  assert (
                      contiguous_inputs[i].dtype == inputs[i].dtype
                  ), f"Dtype mismatch for {i}th input. Expect {inputs[i].dtype}, got {contiguous_inputs[i].dtype}."
  
                  if need_cudagraphs_record:
                      # If cudagraphs is enabled, this memory is reserved for future cudagraph runs
                      # Clone is required to avoid re-using user-provided GPU memory
                      self._input_buffers[i] = contiguous_inputs[i].clone()
                  else:
                      self._input_buffers[i].copy_(contiguous_inputs[i])
  
              if need_cudagraphs_record:
                  # Reconstruct the original args and kwargs structure from static input buffers
                  # using the input specification stored during module compilation
                  args, kwargs = tree_unflatten(
                      self._input_buffers, self.compiled_module._in_spec
                  )
  
              self._caller_stream = torch.cuda.current_stream()
              if (
                  self._engine_stream == torch.cuda.default_stream()
                  or self._engine_stream is None
              ):
                  self._engine_stream = torch.cuda.Stream()
  
              self._engine_stream.wait_stream(self._caller_stream)
  
              with torch.cuda.stream(self._engine_stream):
                  if need_cudagraphs_record:
                      self.cudagraph = torch.cuda.CUDAGraph()
                      with torch.cuda.graph(self.cudagraph, stream=self._engine_stream):
                          self._output_buffers = self.compiled_module(*args, **kwargs)
  
                  self.cudagraph.replay()  # type: ignore
              self._caller_stream.wait_stream(self._engine_stream)
  
              if isinstance(self._output_buffers, (list, tuple)):
                  output_buffers = self._output_buffers
              else:
                  output_buffers = [self._output_buffers]
              outputs = [output.clone() for output in output_buffers]
              if len(outputs) == 1:
                  return outputs[0]
              return outputs
          else:
              self._reset_captured_graph()
              return self.compiled_module(*args, **kwargs)
  