content: import math
  from typing import Dict, Optional, Sequence, Union
  
  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.dynamo._SourceIR import SourceIR
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion.converter_utils import (
      extend_attr_to_tuple,
      get_positive_dim,
      has_dynamic_shape,
      set_layer_name,
  )
  
  
  def avg_poolNd(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      kernel_size: Sequence[int],
      stride: Union[int, Sequence[int]],
      padding: Union[int, Sequence[int]] = 0,
      ceil_mode: bool = False,
      count_include_pad: bool = True,
      divisor_override: Optional[int] = None,
  ) -> TRTTensor:
      if ceil_mode is not False:
          raise RuntimeError("ceil_mode is not yet supported!")
  
      if divisor_override is not None:
          raise RuntimeError("divisor_override is not yet supported!")
  
      dim = len(kernel_size)
  
      kernel_size = extend_attr_to_tuple(kernel_size, dim)
  
      if stride == []:
          stride = kernel_size
      else:
          stride = extend_attr_to_tuple(stride, dim)
  
      padding = extend_attr_to_tuple(padding, dim)
  
      # add average pooling layer
      pool_layer = ctx.net.add_pooling_nd(
          input=input,
          type=trt.PoolingType.AVERAGE,
          window_size=kernel_size,
      )
  
      pool_layer.stride_nd = stride
      pool_layer.padding_nd = padding
      pool_layer.average_count_excludes_padding = not count_include_pad
  
      set_layer_name(pool_layer, target, name, source_ir)
      return pool_layer.get_output(0)
  
  
  def max_poolNd(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      kernel_size: Sequence[int],
      stride: Union[int, Sequence[int]],
      padding: Union[int, Sequence[int]] = 0,
      dilation: Union[int, Sequence[int]] = 1,
      ceil_mode: bool = False,
  ) -> TRTTensor:
      if has_dynamic_shape(input.shape):
          assert input.shape[1] != -1, "Channel dim can't be dynamic for pooling."
  
      if dilation != 1:
          raise RuntimeError("dilation is not yet supported!")
  
      if ceil_mode is not False:
          raise RuntimeError("ceil_mode is not yet supported!")
  
      dim = len(kernel_size)
  
      kernel_size = extend_attr_to_tuple(kernel_size, dim)
  
      if stride == []:
          stride = kernel_size
      else:
          stride = extend_attr_to_tuple(stride, dim)
  
      padding = extend_attr_to_tuple(padding, dim)
  
      # add max pooling layer
      pool_layer = ctx.net.add_pooling_nd(
          input=input,
          type=trt.PoolingType.MAX,
          window_size=kernel_size,
      )
  
      pool_layer.stride_nd = stride
      pool_layer.padding_nd = padding
  
      set_layer_name(pool_layer, target, name, source_ir)
      return pool_layer.get_output(0)
  
  
  def adaptive_avg_pool1d(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      output_size: Union[int, Sequence[int]],
  ) -> TRTTensor:
      def start_index(idx: int, out_dim: int, in_dim: int) -> int:
          """Calculate the start index of each pooling window"""
          return math.floor((float(idx) * float(in_dim)) / out_dim)
  
      def end_index(idx: int, out_dim: int, in_dim: int) -> int:
          """Calculate the end index of each pooling window"""
          return math.ceil((float(idx + 1) * float(in_dim)) / out_dim)
  
      if has_dynamic_shape(input.shape):
          assert (
              input.shape[-1] != -1 and input.shape[-2] != -1
          ), "Last 2 dimensions can't be dynamic for adaptive_avg_pool1d."
  
      in_dim = input.shape[-1]
      out_dim = output_size if isinstance(output_size, int) else output_size[0]
      output_list = []
  
      # store {index: slice} for reducing repeated slice ops
      idx_slice_map: Dict[int, TRTTensor] = {}
      # iterate over each output dimension
      for i in range(out_dim):
          # calculate the start and end index of each pooling window
          start = start_index(i, out_dim, in_dim)
          end = end_index(i, out_dim, in_dim)
  
          # slice the input tensor from start to end index, the result of which is the window waiting for pooling
          slices = []
          for j in range(start, end):
              if j in idx_slice_map:
                  slice = idx_slice_map[j]
              else:
                  slice = impl.select.select(
                      ctx, target, source_ir, f"{name}_select_{j}", input, -1, j
                  )
                  slice = impl.shuffle.reshape(
                      ctx,
                      target,
                      source_ir,
                      f"{name}_reshape_{i}_{j}",
                      slice,
                      (*slice.shape, 1),
                  )
                  idx_slice_map[j] = slice
  
              slices.append(slice)
  
          slices = impl.cat.cat(
              ctx, target, source_ir, f"{name}_slices_cat_{i}", slices, dim=-1
          )
          # calculate the mean of the slices (average pooling output) and append to the output list
          output_list.append(
              impl.reduce.mean(
                  ctx, target, source_ir, f"{name}_sum_{i}", slices, dim=-1, keepdim=True
              )
          )
  
      output = impl.cat.cat(ctx, target, source_ir, f"{name}_cat", output_list, dim=-1)
      return output
  
  
  def adaptive_avg_poolNd(
      ctx: ConversionContext,
      target: Union[Target, str],
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      output_size: Sequence[int],
  ) -> TRTTensor:
      if has_dynamic_shape(input.shape):
          if len(output_size) == 2:  # adaptive_avg_pool2d
              assert (
                  input.shape[-1] != -1 and input.shape[-2] != -1
              ), "Last 2 dimensions can't be dynamic for adaptive_avg_pool2d."
          elif len(output_size) == 3:  # adaptive_avg_pool3d
              assert (
                  input.shape[-1] != -1
                  and input.shape[-2] != -1
                  and input.shape[-3] != -1
              ), "Last 3 dimensions can't be dynamic for adaptive_avg_pool3d."
  
      input_shape = input.shape
      input_rank = len(input_shape)
      output_rank = len(output_size)
      need_reshape_back = False
  
      if input_rank == output_rank + 1:  # reshape to 4D/5D for TRT pooling
          input = impl.shuffle.reshape(
              ctx, target, source_ir, f"{name}_reshape", input, (1, *input.shape)
          )
          need_reshape_back = True
          input_shape = input.shape
          input_rank = len(input_shape)
  
      extend_len = len(output_size)
      output_size = list(output_size)
      original_input = input
  
      # repeat_interleave the input if the dim of output is larger than input
      insert_axises = []
      for axis in range(1, extend_len + 1):
          axis = -axis
          positive_axis = get_positive_dim(
              axis, input_rank
          )  # convert to positive axis, which is for calculating new shapes below
          input_dim = input_shape[axis]
          output_dim = output_size[axis]
          diff = output_dim - input_dim
          if diff > 0:  # the dim of output is larger than input
              times = output_dim // input_dim
              remainder = output_dim % input_dim
              if (
                  diff == 2 and remainder == 2
              ):  # case 1: output_dim - input_dim == 2 and is not an integral multiple
                  insert_axises.append(axis)
                  remainder -= 1
                  output_size[axis] -= 1
  
              if (
                  remainder + 1 == input_dim
              ):  # case 2: remainder + 1 == input_dim, we will repeat_interleave the whole input
                  remainder = 0
                  times += 1
  
              flags = []  # record the axis that needs to be repeated
              concat_list = []
              for j in range(
                  input_dim
              ):  # iterate the input dim to see which dim needs to be repeated or not
                  single_elem = impl.select.select(
                      ctx, target, source_ir, f"{name}_select_{axis}_{j}", input, axis, j
                  )
                  new_shape = list(single_elem.shape)
                  new_shape.insert(positive_axis, 1)
                  single_elem = impl.shuffle.reshape(
                      ctx,
                      target,
                      source_ir,
                      f"{name}_reshape_{axis}_{j}",
                      single_elem,
                      new_shape,
                  )
                  if remainder > 0 or j in flags:
                      concat_list.extend([single_elem] * (times + 1))
                      remainder -= 2
                      flags.append(input_dim - j - 1)
                  else:
                      concat_list.extend([single_elem] * times)
                  out = impl.cat.cat(
                      ctx, target, source_ir, f"{name}_cat_{axis}_{j}", concat_list, axis
                  )
              input = out
  
      stride = tuple(
          input.shape[-extend_len + i] // output_size[i] for i in range(extend_len)
      )
      kernel_size = tuple(
          input.shape[-extend_len + i] - (output_size[i] - 1) * stride[i]
          for i in range(extend_len)
      )
  
      # Don't have to pool, directly return
      if all(s == 1 for s in stride) and all(k == 1 for k in kernel_size):
          if need_reshape_back:  # reshape back
              input = impl.shuffle.reshape(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_reshape_back",
                  input,
                  (*input.shape[1:],),
              )
          return input
  
      layer = ctx.net.add_pooling_nd(
          input=input, type=trt.PoolingType.AVERAGE, window_size=kernel_size
      )
      layer.stride_nd = stride
      set_layer_name(layer, target, f"{name}_pooling_{extend_len}d", source_ir)
  
      output = layer.get_output(0)
  
      # For case 1, we need to split the output and insert the mid of input
      for axis in insert_axises:
          positive_axis = get_positive_dim(axis, input_rank)
          input_dim = input_shape[axis]
          output_dim = output_size[axis]
          if input_dim % 2 == 1:
              prev_one = impl.select.select(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_select_prev_one_{axis}",
                  output,
                  axis,
                  output_dim // 2 - 1,
              )
              extend_shape = list(prev_one.shape)
              extend_shape.insert(positive_axis, 1)
              prev_one = impl.shuffle.reshape(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_reshape_extend_shape_{axis}",
                  prev_one,
                  extend_shape,
              )
              prev_two = impl.select.select(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_select_prev_two_{axis}",
                  output,
                  axis,
                  output_dim // 2 - 2,
              )
              prev_two = impl.shuffle.reshape(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_two_shape_reshape_{axis}",
                  prev_two,
                  extend_shape,
              )
              prev_one_two_diff = impl.elementwise.sub(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_prev_one_two_diff_{axis}",
                  prev_one,
                  prev_two,
              )
  
              mid = impl.elementwise.add(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_mid_{axis}",
                  prev_one,
                  prev_one_two_diff,
              )
              split_output = impl.split.split(
                  ctx, target, source_ir, f"{name}_split_{axis}", output, 2, axis
              )
              split_output.insert(1, mid)
              output = impl.cat.cat(
                  ctx, target, source_ir, f"{name}_cat_{axis}", split_output, axis
              )
          else:
              mid1 = impl.select.select(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_select_{axis}",
                  original_input,
                  axis,
                  input_dim // 2 - 1,
              )
              new_shape = list(mid1.shape)
              new_shape.insert(positive_axis, 1)
              mid1 = impl.shuffle.reshape(
                  ctx, target, source_ir, f"{name}_reshape_{axis}", mid1, new_shape
              )
              mid2 = impl.select.select(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_select_{axis}",
                  original_input,
                  axis,
                  input_dim // 2,
              )
              mid2 = impl.shuffle.reshape(
                  ctx, target, source_ir, f"{name}_reshape_{axis}", mid2, new_shape
              )
              split_output = impl.split.split(
                  ctx,
                  target,
                  source_ir,
                  f"{name}_split_{axis}",
                  output,
                  [output_dim // 2, 1, output_dim // 2],
                  axis,
              )
              split_output[1] = mid1
              split_output.insert(2, mid2)
              output = impl.cat.cat(
                  ctx, target, source_ir, f"{name}_cat_{axis}", split_output, axis
              )
  
      if need_reshape_back:  # reshape back
          output = impl.shuffle.reshape(
              ctx, target, source_ir, f"{name}_reshape_back", output, (*output.shape[1:],)
          )
  
      return output
  