content: # flake8: noqa
  import logging
  import math
  import operator
  import warnings
  from typing import Dict, Optional, Sequence, Tuple, Union, cast
  
  import numpy as np
  
  # @manual=//deeplearning/trt/python:py_tensorrt
  import tensorrt as trt
  import torch
  import torch_tensorrt.fx.tracer.acc_tracer.acc_utils as acc_utils
  from torch.fx.immutable_collections import immutable_list
  from torch.fx.node import Argument, Target
  from torch_tensorrt.fx.converters import acc_ops_converters
  from torch_tensorrt.fx.converters.impl import activation, convolution
  
  from ..converter_registry import tensorrt_converter
  from ..types import *  # noqa: F403
  from .converter_utils import *  # noqa: F403
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  
  ## converter list in alphabetic order
  @tensorrt_converter(torch.ops.aten.add.Tensor)
  def aten_ops_add(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_add(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten._adaptive_avg_pool3d.default)
  @tensorrt_converter(torch.ops.aten._adaptive_avg_pool2d.default)
  def aten_ops_adaptive_avg_poolnd(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "output_size": args[1],
      }
      return acc_ops_converters.acc_ops_adaptive_avg_poolnd(
          network, target, None, kwargs_new, name
      )
  
  
  @tensorrt_converter(torch.ops.aten.mean.default)
  @tensorrt_converter(torch.ops.aten.mean.dim)
  def aten_ops_mean(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> TRTTensor:
      # Default invocation of aten.mean only uses first argument and
      # averages over all elements (all dimensions)
      # aten.mean.dim invocation allows specification of dimensions to average
      # over, as well at the option to keep the dimension or not
      kwargs_new = {
          "input": args[0],
          "dim": args[1] if len(args) >= 2 else list(range(len(args[0].shape))),
          "keepdim": args[2] if len(args) >= 3 else False,
      }
      return add_reduce_layer(
          network, target, args, kwargs_new, trt.ReduceOperation.AVG, name
      )
  
  
  @tensorrt_converter(torch.ops.aten.batch_norm)
  def aten_ops_batch_norm(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "weight": args[1],
          "bias": args[2],
          "running_mean": args[3],
          "running_var": args[4],
          "training": args[5],
          "momentum": args[6],
          "eps": args[7],
      }
      return acc_ops_converters.acc_ops_batch_norm(
          network, target, None, kwargs_new, name
      )
  
  
  @tensorrt_converter(torch.ops.aten.convolution.default)
  def aten_ops_convolution(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "weight": args[1],
          "bias": args[2],
          "stride": args[3],
          "padding": args[4],
          "dilation": args[5],
          "groups": args[8],
      }
      # we do not handle transposed.
      if args[6] is True:
          raise RuntimeError(f"Target {target} does not support `transposed=True` ")
      # we do not handle output_padding.
      if args[7] not in ([0], [0, 0], [0, 0, 0]):
          raise RuntimeError(f"Target {target} has non-0 output_padding")
  
      if len(kwargs_new["stride"]) == 1:
          return convolution.convNd(
              network,
              target,
              source_ir=SourceIR.ATEN,
              name=name,
              is_conv1d=True,
              input_val=kwargs_new["input"],
              weight=kwargs_new["weight"],
              bias=kwargs_new["bias"],
              stride=kwargs_new["stride"],
              padding=kwargs_new["padding"],
              dilation=kwargs_new["dilation"],
              groups=kwargs_new["groups"],
          )
      else:
          return convolution.convNd(
              network,
              target,
              source_ir=SourceIR.ATEN,
              name=name,
              is_conv1d=False,
              input_val=kwargs_new["input"],
              weight=kwargs_new["weight"],
              bias=kwargs_new["bias"],
              stride=kwargs_new["stride"],
              padding=kwargs_new["padding"],
              dilation=kwargs_new["dilation"],
              groups=kwargs_new["groups"],
          )
  
  
  @tensorrt_converter(torch.ops.aten.div.default)
  @tensorrt_converter(torch.ops.aten.div.Tensor_mode)
  @tensorrt_converter(torch.ops.aten.div.Tensor)
  def aten_ops_div(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      rounding_mode = kwargs.get("rounding_mode")
      if rounding_mode is None:
          return acc_ops_converters.acc_ops_div(network, target, None, kwargs_new, name)
      elif rounding_mode == "floor":
          return acc_ops_converters.acc_ops_floor_div(
              network, target, None, kwargs_new, name
          )
      elif rounding_mode == "trunc":
          return acc_ops_converters.acc_ops_trunc_div(
              network, target, None, kwargs_new, name
          )
      else:
          raise RuntimeError(
              f"Target {target} does not support rounding mode {rounding_mode}"
          )
  
  
  @tensorrt_converter(torch.ops.aten.elu.default)
  def aten_ops_elu(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      if len(args) > 2:
          return activation.selu(
              network,
              target,
              SourceIR.ATEN,
              name,
              args[0],
          )
      return activation.elu(
          network,
          target,
          SourceIR.ATEN,
          name,
          args[0],
          args[1],
      )
  
  
  @tensorrt_converter(torch.ops.aten.floor_divide.default)
  def aten_ops_floor_div(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_floor_div(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.fmod.Scalar)
  @tensorrt_converter(torch.ops.aten.fmod.Tensor)
  def aten_ops_fmod(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_fmod(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.hardtanh.default)
  def aten_ops_hardtanh(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      return activation.hardtanh(
          network, target, SourceIR.ATEN, name, args[0], args[1], args[2]
      )
  
  
  @tensorrt_converter(torch.ops.aten.fmod.Tensor)
  def aten_ops_fmod(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_fmod(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.leaky_relu.default)
  def aten_ops_leaky_relu(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
  
  
  @tensorrt_converter(torch.ops.aten.linear)
  def aten_ops_linear(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "weight": args[1],
          "bias": args[2],
      }
  
      return acc_ops_converters.acc_ops_linear(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.max_pool3d)
  @tensorrt_converter(torch.ops.aten.max_pool2d)
  def aten_ops_max_poolnd(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "kernel_size": args[1],
          "stride": (
              args[2]
              if len(args) > 2
              else (None, None) if len(args[1]) == 2 else (None, None, None)
          ),
          "padding": (
              args[3] if len(args) > 3 else (0, 0) if len(args[1]) == 2 else (0, 0, 0)
          ),
          "dilation": (
              args[4] if len(args) > 4 else (1, 1) if len(args[1]) == 2 else (1, 1, 1)
          ),
          "ceil_mode": args[5] if len(args) > 5 else False,
      }
      return acc_ops_converters.acc_ops_max_poolnd(
          network, target, None, kwargs_new, name
      )
  
  
  @tensorrt_converter(torch.ops.aten.mul.Tensor)
  def aten_ops_mul(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_mul(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.pow.Tensor_Scalar)
  @tensorrt_converter(torch.ops.aten.pow.Tensor_Tensor)
  def aten_ops_pow(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "exponent": args[1],
      }
      return acc_ops_converters.acc_ops_pow(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.relu.default)
  def aten_ops_relu(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      return activation.relu(
          network,
          target,
          SourceIR.ATEN,
          name,
          args[0],
      )
  
  
  @tensorrt_converter(torch.ops.aten.sub.Tensor)
  def aten_ops_sub(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_sub(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.view.default)
  def aten_ops_reshape(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      input_val = args[0]
      # for case where input_val is TRTensor
      input_val = get_trt_tensor(network, input_val, f"{name}_input_val")
      shape = args[1]
  
      layer = network.add_shuffle(input_val)
  
      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):
                  trt_shape.append(s)
              else:
                  a = get_trt_tensor(network, s, f"{name}_{i}")
                  trt_shape.append(a)
  
          shape_layer = network.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)
      return layer.get_output(0)
  
  
  @tensorrt_converter(torch.ops.aten.tanh.default)
  def aten_ops_tanh(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      return activation.tanh(
          network,
          target,
          SourceIR.ATEN,
          name,
          args[0],
      )
  
  
  @tensorrt_converter(torch.ops.aten.cat.default)
  def aten_ops_cat(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "tensors": args[0],
          "dim": args[1] if len(args) >= 2 else 0,
      }
      return acc_ops_converters.acc_ops_cat(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.expand.default)
  def aten_ops_expand(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "sizes": args[1],
      }
      return acc_ops_converters.acc_ops_expand_tensor(
          network, target, None, kwargs_new, name
      )
  
  
  @tensorrt_converter(operator.floordiv)
  def aten_ops_operator_floordiv(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_floor_div(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(operator.mul)
  def aten_ops_operator_mul(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_mul(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(operator.add)
  def aten_ops_operator_add(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_add(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(operator.sub)
  def aten_ops_operator_sub(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      kwargs_new = {
          "input": args[0],
          "other": args[1],
      }
      return acc_ops_converters.acc_ops_sub(network, target, None, kwargs_new, name)
  
  
  @tensorrt_converter(torch.ops.aten.sym_numel)
  def aten_ops_sym_numel(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      shape_layer = network.add_shape(args[0])
      set_layer_name(shape_layer, target, "_shape_layer")
      reduce_layer = network.add_reduce(
          shape_layer.get_output(0),
          trt.ReduceOperation.PROD,
          axes=get_axes_for_reduce_op(0, False),
          keep_dims=True,
      )
      set_layer_name(reduce_layer, target, "_reduce_layer")
      return reduce_layer.get_output(0)
  
  
  @tensorrt_converter(torch.ops.aten.sym_size)
  def aten_ops_sym_size(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      shape_layer = network.add_shape(args[0])
      ind = args[1]
      set_layer_name(shape_layer, target, "_shape_layer")
      slice_layer = network.add_slice(
          input=shape_layer.get_output(0),
          start=[ind],
          shape=[1],
          stride=[1],
      )
      set_layer_name(slice_layer, target, "_slice_layer")
      return slice_layer.get_output(0)
  
  
  @tensorrt_converter(torch.ops.aten.sigmoid.default)
  def aten_ops_sigmoid(
      network: TRTNetwork,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[TRTTensor, Sequence[TRTTensor]]:
      return activation.sigmoid(
          network,
          target,
          SourceIR.ATEN,
          name,
          args[0],
      )
  