content: from __future__ import annotations
  
  from typing import List, Optional, Tuple
  
  import numpy as np
  import tensorrt as trt
  import torch
  from tensorrt import ITensor as TRTTensor
  from torch.fx.node import Target
  from torch_tensorrt.dynamo._SourceIR import SourceIR
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion.converter_utils import (
      cast_trt_tensor,
      get_positive_dim,
      get_trt_tensor,
      set_layer_name,
  )
  from torch_tensorrt.dynamo.conversion.impl.elementwise.base import (
      convert_binary_elementwise,
  )
  from torch_tensorrt.dynamo.utils import (
      Frameworks,
      unified_dtype_converter,
  )
  
  
  def shape(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input_val: TRTTensor,
      dim: int,
  ) -> TRTTensor:
      """
      This is the general shape layer implementation in TensorRT.
      sym_size.int ops map to addShape layer in TensorRT and returns
      the dynamic shape of the tensor optionally taking in a dim argument.
      """
      shape_layer = ctx.net.add_shape(input_val)
      input_shape = shape_layer.get_output(0)
      input_shape = cast_trt_tensor(
          ctx,
          input_shape,
          trt.int32,
          name + "_shape_casted",
      )
      set_layer_name(shape_layer, target, name + "_shape", source_ir)
  
      n_dims = len(input_val.shape)
      dim = get_positive_dim(dim, n_dims)
      dim_tensor = get_trt_tensor(ctx, dim, name + "_dim")
      gather_layer = ctx.net.add_gather(input_shape, dim_tensor, axis=0)
      set_layer_name(gather_layer, target, name + "_gather", source_ir)
      input_shape = gather_layer.get_output(0)
  
      return input_shape
  
  
  def get_shape_with_dynamic_shape(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      shape: List[int] | Tuple[int, ...] | torch.Tensor,
      input_val: TRTTensor,
  ) -> TRTTensor:
      """
      Prepare the real output tensor shape for dynamic shape mode tensor input.
      How this functions works:
      Assuming the input_val has actual shape [2048, 256, 512], expected reduce operation
      output shape is [-1, 128, 256], this function should return [2048, 128, 256] as the actual
      reduce operation output shape. Steps of calculations are:
          1. get the actual tensor shape of input_val via add_shape layer;
          2. create a all 0 tensor [0, 0, 0];
          3. run elementwise comparison the [0, 0, 0] and [-1, 128, 256] tensor, get a condition tensor [True, False, False];
          4. use the condition tensor [True, False, False] to do selection between [2048, 256, 512] and [-1, 128, 256], replace
             all -1 dynamic shape dimensions with actual batch_size value;
          5. output shape with actual batch_size as [2048, 128, 256]
  
      Args:
          ctx (ConversionContext): TensorRT ConversionContext object.
          shape: calculated shape of the expected output tensor
          input_val (TRTTensor): A TensorRT ITensor.
          target (Target): Target of fx node.
          name (str): The name we want to assign to the created TensorRT layer.
      Returns:
          TensorRT ITensors that represents the actual shape of the input_val
      """
      # Ger real shape info for input_val
      input_shape = ctx.net.add_shape(input_val).get_output(0)
      input_shape = cast_trt_tensor(
          ctx,
          input_shape,
          trt.int32,
          name + "_int32_casted",
      )
      # input_shape.dtype is int64 in TRT 10.0
      input_np_dtype = unified_dtype_converter(input_shape.dtype, Frameworks.NUMPY)
      scale_layer = ctx.net.add_constant(
          input_shape.shape, np.ascontiguousarray(shape, dtype=input_np_dtype)
      )
      set_layer_name(scale_layer, target, f"{name}_scale")
      scale_res = scale_layer.get_output(0)
  
      length = input_shape.shape[0]
  
      zero_layer = ctx.net.add_constant(
          input_shape.shape, np.zeros((length), dtype=np.int32)
      )
  
      set_layer_name(zero_layer, target, f"{name}_zeros")
  
      condition_val = convert_binary_elementwise(
          ctx,
          target,
          source_ir,
          f"{name}_shape",
          trt.ElementWiseOperation.LESS,
          scale_res,
          zero_layer.get_output(0),
      )
      select_layer = ctx.net.add_select(condition_val, input_shape, scale_res)
      set_layer_name(select_layer, target, f"{name}_select")
      return select_layer.get_output(0)
  