content: from typing import Optional, Sequence, Union
  
  import numpy as np
  import tensorrt as trt
  import torch_tensorrt.dynamo.conversion.impl as impl
  from tensorrt import ITensor as TRTTensor
  from torch.fx.node import Target
  from torch_tensorrt import _enums
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion.converter_utils import (
      SourceIR,
      cast_trt_tensor,
      flatten_dims,
      get_trt_tensor,
      set_layer_name,
  )
  from torch_tensorrt.dynamo.conversion.impl.shape import get_shape_with_dynamic_shape
  
  
  def reshape(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      shape: Sequence[int],
  ) -> TRTTensor:
      layer = ctx.net.add_shuffle(input)
      if all(isinstance(s, int) for s in shape):
          layer.reshape_dims = tuple(shape)
      else:
          # Convert all the dimensions to trt Tensors.
          trt_shape = []
  
          for i, s in enumerate(shape):
              if isinstance(s, TRTTensor):
                  dim_int32 = cast_trt_tensor(
                      ctx,
                      s,
                      _enums.dtype.int32,
                      name + f"_int32_casted_{i}",
                  )
                  trt_shape.append(dim_int32)
              else:
                  a = get_trt_tensor(ctx, s, f"{name}_{i}")
                  trt_shape.append(a)
          shape_layer = ctx.net.add_concatenation(inputs=trt_shape)
          shape_layer.axis = 0
          shape_layer.name = f"{name}_output_shape"
          layer.set_input(1, shape_layer.get_output(0))
  
      set_layer_name(layer, target, name, source_ir)
      return layer.get_output(0)
  
  
  def pixel_shuffle(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      upscale_factor: int,
  ) -> TRTTensor:
      # Get input shape tensor
      input_shape_tensor = get_shape_with_dynamic_shape(
          ctx,
          target,
          source_ir,
          name + "_shape",
          input.shape,
          input,
      )
  
      # Extract in_channels, in_height, and in_width from the input shape tensor
      in_channels_tensor = ctx.net.add_slice(
          input_shape_tensor, start=(len(input.shape) - 3,), shape=(1,), stride=(1,)
      ).get_output(0)
      in_height_tensor = ctx.net.add_slice(
          input_shape_tensor, start=(len(input.shape) - 2,), shape=(1,), stride=(1,)
      ).get_output(0)
      in_width_tensor = ctx.net.add_slice(
          input_shape_tensor, start=(len(input.shape) - 1,), shape=(1,), stride=(1,)
      ).get_output(0)
  
      # Calculate out_channels, out_height, and out_width as tensors
      upscale_factor_sq = upscale_factor * upscale_factor
      upscale_factor_tensor = get_trt_tensor(
          ctx, upscale_factor, f"{name}_upscale_factor"
      )
      upscale_factor_sq_tensor = get_trt_tensor(
          ctx, upscale_factor_sq, f"{name}_upscale_factor_sq"
      )
  
      out_channels_tensor = impl.elementwise.floor_divide(
          ctx,
          target,
          source_ir,
          f"{name}_out_channels_tensor",
          in_channels_tensor,
          upscale_factor_sq_tensor,
      )
      out_height_tensor = impl.elementwise.mul(
          ctx,
          target,
          source_ir,
          f"{name}_out_height_tensor",
          in_height_tensor,
          upscale_factor,
      )
      out_width_tensor = impl.elementwise.mul(
          ctx,
          target,
          source_ir,
          f"{name}_out_width_tensor",
          in_width_tensor,
          upscale_factor,
      )
  
      # Construct new shape tensor
      new_shape_tensors = [
          ctx.net.add_slice(
              input_shape_tensor, start=(i,), shape=(1,), stride=(1,)
          ).get_output(0)
          for i in range(len(input.shape) - 3)
      ]
      new_shape_tensors += [
          out_channels_tensor,
          upscale_factor_tensor,
          upscale_factor_tensor,
          in_height_tensor,
          in_width_tensor,
      ]
  
      # Reshape tensor
      reshaped_tensor = reshape(
          ctx, target, source_ir, f"{name}_reshape", input, new_shape_tensors
      )
  
      # Permute shape
      rank = len(input.shape)
      permute_shape = list(range(rank))
      permute_shape.insert(-2, rank)
      permute_shape.insert(-1, rank + 1)
      permuted_tensor = impl.permutation.permute(
          ctx, target, source_ir, f"{name}_permute", reshaped_tensor, permute_shape
      )
  
      # Construct output shape tensor
      out_shape_tensors = [
          ctx.net.add_slice(
              input_shape_tensor, start=(i,), shape=(1,), stride=(1,)
          ).get_output(0)
          for i in range(len(input.shape) - 3)
      ]
      out_shape_tensors += [out_channels_tensor, out_height_tensor, out_width_tensor]
  
      return reshape(
          ctx,
          target,
          source_ir,
          f"{name}_reshape_out",
          permuted_tensor,
          out_shape_tensors,
      )
  
  
  def pixel_unshuffle(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      downscale_factor: int,
  ) -> TRTTensor:
      # Get input shape tensor
      input_shape_tensor = get_shape_with_dynamic_shape(
          ctx,
          target,
          source_ir,
          name + "_shape",
          input.shape,
          input,
      )
  
      # Extract in_channels, in_height, and in_width from the input shape tensor
      in_channels_tensor = ctx.net.add_slice(
          input_shape_tensor, start=(len(input.shape) - 3,), shape=(1,), stride=(1,)
      ).get_output(0)
      in_height_tensor = ctx.net.add_slice(
          input_shape_tensor, start=(len(input.shape) - 2,), shape=(1,), stride=(1,)
      ).get_output(0)
      in_width_tensor = ctx.net.add_slice(
          input_shape_tensor, start=(len(input.shape) - 1,), shape=(1,), stride=(1,)
      ).get_output(0)
  
      # Calculate out_channels, out_height, and out_width as tensors
      downscale_factor_sq = downscale_factor * downscale_factor
      downscale_factor_tensor = get_trt_tensor(
          ctx, downscale_factor, f"{name}_downscale_factor"
      )
      downscale_factor_sq_tensor = get_trt_tensor(
          ctx, downscale_factor_sq, f"{name}_downscale_factor_sq"
      )
  
      out_channels_tensor = impl.elementwise.mul(
          ctx,
          target,
          source_ir,
          f"{name}_out_channels_tensor",
          in_channels_tensor,
          downscale_factor_sq_tensor,
      )
      out_height_tensor = impl.elementwise.floor_divide(
          ctx,
          target,
          source_ir,
          f"{name}_out_height_tensor",
          in_height_tensor,
          downscale_factor_tensor,
      )
      out_width_tensor = impl.elementwise.floor_divide(
          ctx,
          target,
          source_ir,
          f"{name}_out_width_tensor",
          in_width_tensor,
          downscale_factor_tensor,
      )
  
      # Construct new shape tensor
      new_shape_tensors = [
          ctx.net.add_slice(
              input_shape_tensor, start=(i,), shape=(1,), stride=(1,)
          ).get_output(0)
          for i in range(len(input.shape) - 3)
      ]
      new_shape_tensors += [
          in_channels_tensor,
          out_height_tensor,
          downscale_factor_tensor,
          out_width_tensor,
          downscale_factor_tensor,
      ]
  
      reshaped_tensor = reshape(
          ctx, target, source_ir, f"{name}_reshape", input, new_shape_tensors
      )
  
      # Permute shape
      rank = len(new_shape_tensors)
      permute_shape = list(range(rank - 5)) + [
          rank - 5,  # in_channels
          rank - 3,  # downscale_factor
          rank - 1,  # downscale_factor
          rank - 4,  # out_height
          rank - 2,  # out_width
      ]
      permuted_tensor = impl.permutation.permute(
          ctx, target, source_ir, f"{name}_permute", reshaped_tensor, permute_shape
      )
  
      # Construct output shape tensor
      out_shape_tensors = [
          ctx.net.add_slice(
              input_shape_tensor, start=(i,), shape=(1,), stride=(1,)
          ).get_output(0)
          for i in range(len(input.shape) - 3)
      ]
      out_shape_tensors += [out_channels_tensor, out_height_tensor, out_width_tensor]
  
      return reshape(
          ctx,
          target,
          source_ir,
          f"{name}_reshape_out",
          permuted_tensor,
          out_shape_tensors,
      )
  
  
  def resize(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      sizes: Sequence[int],
  ) -> TRTTensor:
      input_np_dtype = _enums.dtype._from(input.dtype).to(np.dtype)
      input_val = get_trt_tensor(ctx, input, f"{name}_input")
  
      # Calculate the total number of elements for new and current shape
      new_num_elements = np.prod(sizes)
      current_num_elements = np.prod(input_val.shape)
  
      if new_num_elements > current_num_elements:
          # Create a padding tensor with the required size and initialize new elements with zeros
          padding_size = new_num_elements - current_num_elements
          padding_tensor = ctx.net.add_constant(
              (padding_size,), trt.Weights(np.zeros(padding_size, dtype=input_np_dtype))
          ).get_output(0)
  
          # Flatten input tensor to 1D for concatenation
          flatten_shape = flatten_dims(input_val, 0, -1)
          flattened_input = reshape(
              ctx, target, source_ir, f"{name}_flatten_input", input_val, flatten_shape
          )
  
          # Concatenate the flattened input tensor and padding tensor
          reshaped_tensor = impl.cat.cat(
              ctx,
              target,
              source_ir,
              f"{name}_cat",
              [flattened_input, padding_tensor],
              dim=0,
          )
      elif new_num_elements < current_num_elements:
          # Flatten input tensor to 1D for slicing
          flatten_shape = flatten_dims(input_val, 0, -1)
          flattened_input = reshape(
              ctx, target, source_ir, f"{name}_flatten_input", input_val, flatten_shape
          )
  
          # Slice the flattened input tensor to the desired number of elements
          slice_layer = ctx.net.add_slice(flattened_input, [0], [new_num_elements], [1])
          reshaped_tensor = slice_layer.get_output(0)
      else:
          reshaped_tensor = input_val
  
      # Reshape the final output tensor to the target sizes
      resized_output = reshape(
          ctx, target, source_ir, f"{name}_final_reshape", reshaped_tensor, sizes
      )
  
      return resized_output
  