content: from typing import Optional, Union
  
  import numpy as np
  import tensorrt as trt
  import torch
  from tensorrt import ITensor as TRTTensor
  from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
  from torch.fx.node import Target
  from torch_tensorrt import ENABLED_FEATURES
  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 (
      get_trt_tensor,
      set_layer_name,
      to_torch,
  )
  
  
  def get_ir(target: Target) -> SourceIR:
      target_module = getattr(target, "__module__", "None")
      if any(
          target_module.startswith(prefix)
          for prefix in ("torch.ops.aten", "torch._ops.aten")
      ):
          return SourceIR.ATEN
      elif any(
          target_module.startswith(prefix)
          for prefix in ("torch.ops.prims", "torch._ops.prims")
      ):
          return SourceIR.PRIM
      elif target_module.startswith("torch.nn"):
          return SourceIR.NN
  
      return SourceIR.UNKNOWN
  
  
  def validate_int8_activation_quantization(name: str, dtype: trt.DataType) -> None:
      if (
          dtype == trt.DataType.INT8
          and ".input_quantizer" in name
          and ENABLED_FEATURES.tensorrt_rtx
      ):
          # RTX does not support int8 activation quantization
          # TODO: lan to remove this once rtx team has added the support for int8 activation quantization
          raise NotImplementedError(
              "TensorRT-RTX does not support int8 activation quantization, only support int8 weight quantization"
          )
  
  
  def quantize(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input_tensor: TRTTensor,
      amax: Union[np.ndarray, torch.Tensor],
      num_bits: int,
      exponent_bits: int,
  ) -> TRTTensor:
      """
      Adds quantize and dequantize ops (QDQ) which quantize to INT8 or FP8 based
      on the output_type set and dequantizes them back.
      """
      with unset_fake_temporarily():
          if isinstance(input_tensor, (torch.Tensor, TRTTensor)):
              if input_tensor.dtype not in (
                  trt.float32,
                  trt.float16,
                  trt.bfloat16,
                  torch.bfloat16,
                  torch.float16,
                  torch.float32,
              ):
                  raise ValueError(
                      f"quantize converter received an input of {input_tensor.dtype} type. Supported types: float32 | float16 | bfloat16"
                  )
              if num_bits != 8 or exponent_bits not in (0, 4):
                  raise ValueError(
                      f"quantize converter currently only accept INT8 or FP8 based quantize, got {num_bits=}, {exponent_bits=}"
                  )
          else:
              raise ValueError(
                  f"quantize converter received an input of {type(input_tensor)} type. Supported types: torch.Tensor | TRTTensor"
              )
  
          if num_bits == 8 and exponent_bits == 0:
              dtype = trt.DataType.INT8
              max_bound = 127
          elif num_bits == 8 and exponent_bits == 4:
              dtype = trt.DataType.FP8
              max_bound = 448
  
          validate_int8_activation_quantization(name, dtype)
  
          axis = None
          # int8 weight quantization is per-channel quantization(it can have one or multiple amax values)
          if dtype == trt.DataType.INT8 and amax.numel() > 1:
              # if the amax has more than one element, calculate the axis, otherwise axis value will be ignored
              amax_init_shape = amax.shape
              amax = amax.squeeze().data
              assert (
                  len(amax.shape) == 1
              ), f"TensorRT does not support multi-axis quantization. {name=} {amax_init_shape=} {amax.shape=} "
              axis = list(amax_init_shape).index(list(amax.shape)[0])
              assert (
                  axis == 0
              ), f"{name=} {amax=} is per-channel quantization, expected axis to be 0, but got {axis=}"
          else:
              # int8 activation and fp8 weight/activation quantization is per-tensor quantization, it can only have single amax value
              assert (
                  amax.numel() == 1
              ), f"{name=} is per-tensor quantization, expected amax is a singular value, but got {amax.shape=}"
  
          if not isinstance(amax, trt.ITensor):
              amax = to_torch(amax, None)
              scale = torch.divide(amax, max_bound)
              scale = get_trt_tensor(ctx, scale, name + "_scale", dtype=torch.float32)
          else:
              scale = impl.elementwise.div(
                  ctx,
                  target,
                  get_ir(target),
                  name,
                  amax,
                  max_bound,
              )
              scale = get_trt_tensor(ctx, scale, name + "_scale", dtype=torch.float32)
  
          # Add Q node
          if num_bits == 8 and exponent_bits == 0:
              dtype = trt.DataType.INT8
          elif num_bits == 8 and exponent_bits == 4:
              dtype = trt.DataType.FP8
  
          if not isinstance(input_tensor, TRTTensor):
              input_tensor = get_trt_tensor(ctx, input_tensor, name + "_quantize_input")
  
          # Add Q node
          quantize_layer = ctx.net.add_quantize(input_tensor, scale, dtype)
          if axis is not None:
              quantize_layer.axis = axis
          set_layer_name(quantize_layer, target, name + "_quantize", source_ir)
          q_output = quantize_layer.get_output(0)
          # Add DQ node
          dequantize_layer = ctx.net.add_dequantize(
              q_output, scale, output_type=input_tensor.dtype
          )
          if axis is not None:
              dequantize_layer.axis = axis
          set_layer_name(dequantize_layer, target, name + "_dequantize", source_ir)
          dq_output = dequantize_layer.get_output(0)
  
          return dq_output
  