content: from typing import Optional, Tuple, Union
  
  import tensorrt as trt
  from tensorrt import ITensor as TRTTensor
  from torch.fx.node import Target
  from torch_tensorrt.dynamo._SourceIR import SourceIR
  from torch_tensorrt.dynamo.conversion import impl
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion.converter_utils import (
      cast_trt_tensor,
      flatten_dims,
      get_axes_for_reduce_op,
      get_positive_dim,
      get_trt_tensor,
      has_dynamic_shape,
      set_layer_name,
  )
  from torch_tensorrt.dynamo.conversion.impl.elementwise import convert_binary_elementwise
  from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape
  from torch_tensorrt.dynamo.utils import DYNAMIC_DIM
  
  
  def argmax_argmin(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      topk_option: trt.TopKOperation,
      dim: Optional[int],
      keep_dim: bool = False,
  ) -> TRTTensor:
      if input.dtype == trt.int32:
          input = cast_trt_tensor(ctx, input, trt.float32, name, target, source_ir)
  
      # Three different cases here:
      # 1. dim == None, flatten input tensor first, keep_dim will be ignore and the output rank == input rank
      # 2. input rank == 1: TopK layer does not support 1 dimensional topk operation. Broadcast input to rank == 2
      # 3. normal cases, no additional handlings
      out = input
      is_dynamic_present = has_dynamic_shape(input.shape)
  
      if dim is None:
          if is_dynamic_present and len(input.shape) != 1:
              multiplier = get_trt_tensor(ctx, 1, name + "_shape")
              for i in range(0, len(input.shape)):
                  if input.shape[i] != DYNAMIC_DIM:
                      multiplier = convert_binary_elementwise(
                          ctx,
                          target,
                          source_ir,
                          name + f"_shape_{i}",
                          trt.ElementWiseOperation.PROD,
                          multiplier,
                          input.shape[i],
                      )
                  else:
                      multiplier = convert_binary_elementwise(
                          ctx,
                          target,
                          source_ir,
                          name + f"_shape_{i}",
                          trt.ElementWiseOperation.PROD,
                          multiplier,
                          get_shape(
                              ctx,
                              target,
                              source_ir,
                              name + f"_shape_dim_stop_{i}",
                              input,
                              i,
                          ),
                      )
              # form shape tensor
              new_shape_layer = ctx.net.add_concatenation(
                  [multiplier, get_trt_tensor(ctx, 1, name + "_one_shape")]
              )
              set_layer_name(
                  new_shape_layer, target, name + "_new_shape_concat", source_ir
              )
              concat_tensor = new_shape_layer.get_output(0)
  
              reshape_dynamic_layer = ctx.net.add_shuffle(input)
              reshape_dynamic_layer.set_input(1, concat_tensor)
              set_layer_name(
                  reshape_dynamic_layer, target, name + "_reshape_layer", source_ir
              )
              out = reshape_dynamic_layer.get_output(0)
  
          else:
              new_shape = (*flatten_dims(input, 0, -1), 1)
              out = impl.shuffle.reshape(
                  ctx, target, source_ir, f"{name}_flatten", input, new_shape
              )
      elif len(input.shape) == 1:
          new_shape = (*input.shape, 1)
          out = impl.shuffle.reshape(
              ctx, target, source_ir, f"{name}_broadcast", input, new_shape
          )
  
      # Reduce over the flattened input if the dimension is None, otherwise the specified dimension
      reduce_mask = get_axes_for_reduce_op(
          get_positive_dim(dim if dim is not None else 0, len(out.shape))
      )
  
      topk_layer = ctx.net.add_topk(out, topk_option, 1, reduce_mask)
      set_layer_name(topk_layer, target, name, source_ir)
  
      out = topk_layer.get_output(1)
  
      if dim is None:
          new_shape = ((1,) * len(input.shape)) if keep_dim else ()
          out = impl.shuffle.reshape(
              ctx, target, source_ir, f"{name}_unflatten", out, new_shape
          )
      elif len(input.shape) == 1:
          out = impl.squeeze.squeeze(
              ctx,
              target,
              source_ir,
              f"{name}_squeeze",
              out,
              1 if keep_dim else (0, 1),
          )
      elif not keep_dim:
          out = impl.squeeze.squeeze(ctx, target, source_ir, f"{name}_squeeze", out, dim)
  
      return out
  
  
  def argmax(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: Optional[int],
      keep_dim: bool = False,
  ) -> TRTTensor:
      return argmax_argmin(
          ctx, target, source_ir, name, input, trt.TopKOperation.MAX, dim, keep_dim
      )
  
  
  def argmin(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: Optional[int],
      keep_dim: bool = False,
  ) -> TRTTensor:
      return argmax_argmin(
          ctx, target, source_ir, name, input, trt.TopKOperation.MIN, dim, keep_dim
      )
  
  
  def sort(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      dim: int,
      descending: bool,
      return_indices: bool = True,
  ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]:
      dim = get_positive_dim(dim, len(input.shape))
      k = input.shape[dim]
      return topk(
          ctx,
          target,
          source_ir,
          name,
          input,
          k,
          dim,
          descending,
          sorted=None,
          return_indices=return_indices,
      )
  
  
  def topk(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      k: int,
      dim: int,
      largest: bool,
      sorted: Optional[bool],
      return_indices: bool = True,
  ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]:
      if largest:
          topk_layer = ctx.net.add_topk(
              input,
              trt.TopKOperation.MAX,
              k,
              get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape))),
          )
      else:
          topk_layer = ctx.net.add_topk(
              input,
              trt.TopKOperation.MIN,
              k,
              get_axes_for_reduce_op(get_positive_dim(dim, len(input.shape))),
          )
  
      # TensorRT ITopKLayer does not have a sorted flag, it is always returning the sorted topk elements
      # so here no matter sorted is True or False the returned the topk Tensor object is always sorted
      set_layer_name(topk_layer, target, f"{name}_topk", source_ir)
  
      if return_indices:
          return topk_layer.get_output(0), topk_layer.get_output(1)
      else:
          return topk_layer.get_output(0)
  