content: import logging
  import time
  import unittest
  from typing import Callable, List, Optional, Set, Tuple
  
  # @manual=//deeplearning/trt/python:py_tensorrt
  import tensorrt as trt
  import torch
  import torch.fx
  import torch_tensorrt.fx.tracer.acc_tracer.acc_tracer as acc_tracer
  import torch_tensorrt.fx.tracer.dispatch_tracer.aten_tracer as aten_tracer
  from torch.fx.experimental.normalize import NormalizeArgs
  from torch.fx.passes import shape_prop
  from torch.fx.passes.infra.pass_base import PassResult
  from torch.testing._internal.common_utils import TestCase
  from torch_tensorrt.fx import InputTensorSpec, TRTInterpreter, TRTModule
  from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
      compose_bmm,
      compose_chunk,
      compose_getitem_slice,
      remove_ops,
      replace_aten_op_with_indices,
      replace_aten_reshape_alias_with_replace,
      replace_builtin_ops,
      replace_native_layernorm_with_layernorm,
      replace_transpose_mm_op_with_linear,
      run_const_fold,
  )
  from torch_tensorrt.fx.passes.pass_utils import chain_passes
  from torch_tensorrt.fx.utils import LowerPrecision, proxytensor_trace
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  
  def fetch_attr(mod, target):
      """
      Fetch an attribute from the ``Module`` hierarchy of ``mod.module``.
  
      Args:
          target (str): The fully-qualfiied name of the attribute to fetch
  
      Return:
          Any: The value of the attribute.
      """
      target_atoms = target.split(".")
      attr_itr = mod
      for i, atom in enumerate(target_atoms):
          if not hasattr(attr_itr, atom):
              raise RuntimeError(
                  f"Node referenced nonexistent target {'.'.join(target_atoms[:i])}"
              )
          attr_itr = getattr(attr_itr, atom)
      return attr_itr
  
  
  @unittest.skipIf(not torch.cuda.is_available(), "Skip because CUDA is not available")
  class TRTTestCase(TestCase):
      def setUp(self):
          super().setUp()
          torch.manual_seed(3)
  
      def run_test(
          self,
          mod,
          inputs,
          expected_ops,
          unexpected_ops,
          interpreter,
          rtol,
          atol,
          precision=LowerPrecision.FP32,
      ):
          with torch.no_grad():
              cuda_inputs = []
              for i in inputs:
                  cuda_inputs.append(i.cuda())
  
              mod.eval()
              if len(expected_ops):
                  self.assert_has_op(mod, expected_ops)
              if unexpected_ops:
                  self.assert_unexpected_op(mod, unexpected_ops)
              start = time.perf_counter()
              interpreter_result = interpreter.run(lower_precision=precision)
              sec = time.perf_counter() - start
              _LOGGER.info(f"Interpreter run time(s): {sec}")
              trt_mod = TRTModule(
                  interpreter_result.engine,
                  interpreter_result.input_names,
                  interpreter_result.output_names,
              )
  
              ref_outputs = mod(*inputs)
  
              torch.cuda.synchronize()
              start_event = torch.cuda.Event(enable_timing=True)
              end_event = torch.cuda.Event(enable_timing=True)
              start_event.record()
              outputs = trt_mod(*cuda_inputs)
              end_event.record()
              torch.cuda.synchronize()
              _LOGGER.info(
                  f"TRT run time(s)= {(start_event.elapsed_time(end_event) * 1.0e-3)}"
              )
  
              if type(outputs) not in (list, tuple):
                  outputs = [outputs]
              if type(ref_outputs) not in (
                  list,
                  tuple,
                  torch.return_types.max,
                  torch.return_types.min,
              ):
                  ref_outputs = [ref_outputs]
              for out, ref in zip(outputs, ref_outputs):
                  if not isinstance(ref, torch.Tensor):
                      ref = torch.tensor([ref])
                  ref = ref.cpu()  # to_dtype test has cases with gpu output
                  if ref.dtype == torch.int64:
                      ref = ref.int()  # convert torch.max's index output tensor to int32
                  torch.testing.assert_close(
                      out.cpu(), ref, rtol=rtol, atol=atol, equal_nan=True
                  )
  
      def run_test_custom_compare_results(
          self,
          mod,
          inputs,
          expected_ops,
          interpreter,
          comparators: List[Tuple[Callable, List]],
          fp16_mode=False,
      ):
          """
          Runs the test and compares the result using the provided comparators.
          The size of comparators must be equal to the number of outputs from 'mod'.
  
          mod          - a model to run.
          inputs       - a list of the model inputs.
          expected ops - a list of ops that should be verified.
          interpreter  - used for converting the model to TRT.
          comparators  - a list of (func, args) pairs corresponding to each of
                         the module outputs. usage: func(x, y, *args)
  
          """
          with torch.no_grad():
              cuda_inputs = []
              for i in inputs:
                  cuda_inputs.append(i.cuda())
  
              mod.eval()
              if len(expected_ops):
                  self.assert_has_op(mod, expected_ops)
  
              interpreter_result = interpreter.run(
                  lower_precision=(
                      LowerPrecision.FP16 if fp16_mode else LowerPrecision.FP32
                  )
              )
              trt_mod = TRTModule(
                  interpreter_result.engine,
                  interpreter_result.input_names,
                  interpreter_result.output_names,
              )
              res_trt = trt_mod(*cuda_inputs).cpu()
              res_cpu = mod(*inputs)
              assert len(res_trt) == len(res_cpu)
              assert len(res_cpu) == len(comparators)
              for output_trt, output_cpu, comparator in zip(
                  res_trt, res_cpu, comparators
              ):
                  comp_func = comparator[0]
                  args = comparator[1]
                  self.assertTrue(comp_func(output_trt, output_cpu, *args))
  
      def run_test_with_error(self, mod, inputs, interpreter, expect_error):
          with self.assertRaises(expect_error):
              with torch.no_grad():
                  cuda_inputs = []
                  for i in inputs:
                      cuda_inputs.append(i.cuda())
  
                  mod.eval()
                  interpreter.run(lower_precision=LowerPrecision.FP32)
  
      def assert_has_op(self, mod, ops):
          ops_in_mod = set()
  
          for node in mod.graph.nodes:
              if node.op == "call_module":
                  ops_in_mod.add(type(fetch_attr(mod, node.target)))
              elif node.op in {"call_function", "call_method"}:
                  ops_in_mod.add(node.target)
  
          self.assertTrue(
              ops_in_mod >= ops, f"expected ops {ops}, actuall ops {ops_in_mod}"
          )
  
      def assert_unexpected_op(self, mod, ops):
          for node in mod.graph.nodes:
              if node.op == "call_module":
                  if type(fetch_attr(mod, node.target)) in ops:
                      return False
              elif node.op in {"call_function", "call_method"}:
                  if node.target in ops:
                      return False
          return True
  
  
  class VanillaTestCase(TRTTestCase):
      def run_test(self, mod, inputs, expected_ops, rtol=1e-03, atol=1e-03):
          mod = torch.fx.symbolic_trace(mod)
          shape_prop.ShapeProp(mod).propagate(*inputs)
          mod = NormalizeArgs(mod).transform()
          interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
          super().run_test(mod, inputs, expected_ops, None, interp, rtol, atol)
  
      def run_test_custom_compare_results(
          self,
          mod,
          inputs,
          expected_ops,
          interpreter,
          comparators: List[Tuple[Callable, List]],
          fp16_mode=False,
      ):
          # interpreter is ignored, we do not need this for Vanilla tests
          # Note this is different from internal version, we need to fix the test case
          # after we refactor the internal callsites to use this file
          mod = torch.fx.symbolic_trace(mod)
          shape_prop.ShapeProp(mod).propagate(*inputs)
          mod = NormalizeArgs(mod).transform()
          interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
          super().run_test_custom_compare_results(
              mod, inputs, expected_ops, interp, comparators, fp16_mode=fp16_mode
          )
  
  
  class AccTestCase(TRTTestCase):
      def run_test(
          self,
          mod,
          inputs,
          expected_ops,
          unexpected_ops=None,
          apply_passes=None,
          test_explicit_batch_dim=True,
          test_implicit_batch_dim=True,
          test_explicit_precision=False,
          rtol=1e-03,
          atol=1e-03,
          precision=LowerPrecision.FP32,
      ):
          mod.eval()
          mod = acc_tracer.trace(mod, inputs)
  
          if apply_passes is not None:
              pass_tracer = chain_passes(*apply_passes)
              mod = pass_tracer(mod, inputs)
  
          if trt.__version__ >= "8.6":
              test_implicit_batch_dim = False
          if test_implicit_batch_dim:
              interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
              )
  
          if test_explicit_batch_dim:
              interp = TRTInterpreter(
                  mod, InputTensorSpec.from_tensors(inputs), explicit_batch_dimension=True
              )
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
              )
  
          if test_explicit_precision:
              interp = TRTInterpreter(
                  mod,
                  InputTensorSpec.from_tensors(inputs),
                  explicit_precision=test_explicit_precision,
              )
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol
              )
  
              interp = TRTInterpreter(
                  mod,
                  InputTensorSpec.from_tensors(inputs),
                  explicit_batch_dimension=True,
                  explicit_precision=test_explicit_precision,
              )
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
              )
  
      def run_test_with_assert_error(
          self,
          mod,
          inputs,
          expect_error,
          test_explicit_batch_dim=True,
          test_implicit_batch_dim=True,
      ):
          mod.eval()
          mod = acc_tracer.trace(mod, inputs)
  
          if test_implicit_batch_dim:
              interp = TRTInterpreter(mod, InputTensorSpec.from_tensors(inputs))
              super().run_test_with_error(mod, inputs, interp, expect_error)
  
          if test_explicit_batch_dim:
              interp = TRTInterpreter(
                  mod, InputTensorSpec.from_tensors(inputs), explicit_batch_dimension=True
              )
              super().run_test_with_error(mod, inputs, interp, expect_error)
  
      def run_test_with_dynamic_shape(
          self,
          mod,
          input_specs,
          expected_ops,
          unexpected_ops=None,
          rtol=1e-03,
          atol=1e-03,
      ):
          mod.eval()
          inputs = InputTensorSpec.create_inputs_from_specs(input_specs)
          mod = acc_tracer.trace(mod, inputs)
          interp = TRTInterpreter(mod, input_specs, explicit_batch_dimension=True)
          super().run_test(mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol)
  
  
  class DispatchTestCase(TRTTestCase):
      def generate_graph(
          self,
          mod: torch.nn.Module,
          original_inputs: List[torch.Tensor],
          expected_ops: Set[Callable],
          unexpected_ops: Optional[Set[Callable]] = None,
          customized_passes: List[Callable] = None,
      ):
          # Torchdynamo+aot proxytensor tracer
          # Below are common passes
          passes_list = [
              compose_bmm,
              compose_chunk,
              compose_getitem_slice,
              replace_aten_reshape_alias_with_replace,
              replace_aten_op_with_indices,
              replace_transpose_mm_op_with_linear,  # after compose_bmm
              replace_native_layernorm_with_layernorm,
              remove_ops,
              replace_builtin_ops,  # after replace_native_layernorm_with_layernorm
          ]
          # Combine with customized passes specific to any model
          if customized_passes:
              passes_list.extend(customized_passes)
          fx_module, _ = aten_tracer.trace(mod, original_inputs)
          for passes in passes_list:
              pr: PassResult = passes(fx_module)
              fx_module = pr.graph_module
          fx_module(*original_inputs)
  
          fx_module = run_const_fold(fx_module)
          _LOGGER.info(f"FX graph= {fx_module.graph}")
  
          if len(expected_ops):
              self.assert_has_op(fx_module, expected_ops)
          if unexpected_ops:
              self.assert_unexpected_op(fx_module, unexpected_ops)
  
          return fx_module
  
      def run_test(
          self,
          mod,
          inputs,
          expected_ops,
          unexpected_ops=None,
          apply_passes=None,
          test_explicit_batch_dim=True,
          test_explicit_precision=False,
          rtol=1e-03,
          atol=1e-03,
          precision=LowerPrecision.FP32,
      ):
          mod.eval()
          mod = self.generate_graph(mod, inputs, expected_ops, unexpected_ops, None)
  
          if apply_passes is not None:
              pass_tracer = chain_passes(*apply_passes)
              mod = pass_tracer(mod, inputs)
  
          if test_explicit_batch_dim:
              interp = TRTInterpreter(
                  mod,
                  InputTensorSpec.from_tensors(inputs),
                  explicit_batch_dimension=True,
              )
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
              )
  
          if test_explicit_precision:
              interp = TRTInterpreter(
                  mod,
                  InputTensorSpec.from_tensors(inputs),
                  explicit_precision=test_explicit_precision,
              )
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol
              )
  
              interp = TRTInterpreter(
                  mod,
                  InputTensorSpec.from_tensors(inputs),
                  explicit_batch_dimension=True,
                  explicit_precision=test_explicit_precision,
              )
              super().run_test(
                  mod, inputs, expected_ops, unexpected_ops, interp, rtol, atol, precision
              )
  
      def run_test_with_dynamic_shape(
          self,
          mod,
          input_specs,
          expected_ops,
          unexpected_ops=None,
          rtol=1e-03,
          atol=1e-03,
      ):
          mod.eval()
          inputs = InputTensorSpec.create_inputs_from_specs(input_specs)
          mod = self.generate_graph(mod, inputs, expected_ops, unexpected_ops, None)
  
          interp = TRTInterpreter(
              mod,
              input_specs,
              explicit_batch_dimension=True,
          )
          # Since the lowering is based on optimal shape. We need to test with
          # different shape(for ex. max shape) for testing dynamic shape
          inputs_max = InputTensorSpec.create_inputs_from_max_specs(input_specs)
          super().run_test(
              mod, inputs_max, expected_ops, unexpected_ops, interp, rtol, atol
          )
  