content: # Owner(s): ["oncall: fx"]
  import logging
  import operator
  import unittest
  from typing import Callable, Dict, List, NamedTuple, Optional, Tuple
  
  import numpy as np
  import torch
  import torch.nn as nn
  
  import torch_tensorrt.fx.tracer.acc_tracer.acc_normalizer as acc_normalizer
  import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
  import torch_tensorrt.fx.tracer.acc_tracer.acc_tracer as acc_tracer
  import torch_tensorrt.fx.tracer.acc_tracer.acc_utils as acc_utils
  import torchvision
  from parameterized import param, parameterized
  
  torch.manual_seed(0)
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  torch.fx.wrap("len")
  
  
  class AccTracerTest(unittest.TestCase):
      def _make_model_unit_test(
          self,
          model,
          *args,
          input_shape=None,
          enable_allclose=False,
          **kwargs,
      ):
          """
          Test that the model can be traced correctly and is producing correct
          result.
          """
          if input_shape is None:
              input_shape = [1, 3, 224, 224]
          input = torch.randn(input_shape)
          traced = acc_tracer.trace(model, [input])
          if enable_allclose:
              torch.testing.assert_close(model(input), traced(input))
          else:
              self.assertTrue(torch.equal(model(input), traced(input)))
          traced_again = acc_tracer.trace(traced, [input])
          if enable_allclose:
              torch.testing.assert_close(model(input), traced_again(input))
          else:
              self.assertTrue(torch.equal(model(input), traced_again(input)))
  
      def _make_acc_op_function_test(
          self,
          acc_op: Callable,
          torch_op,
          *args,
          input_shape=(2, 3),
          validate_same_kwargs=True,
          enable_allclose=False,
          **kwargs,
      ):
          """
          Test that acc_op is traced somewhat.
          """
  
          class TestModule(torch.nn.Module):
              def __init__(self, torch_op, args, kwargs):
                  super().__init__()
                  self._torch_op = torch_op
                  self._args = args
                  self._kwargs = kwargs
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self._torch_op(a, *self._args, **self._kwargs)
  
          m = TestModule(torch_op, args, kwargs)
          m.eval()
          a = torch.randn(*input_shape)
          traced = acc_tracer.trace(m, [a])
          ph_a = acc_op_node = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  if str(node.target) == "a":
                      ph_a = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_op)
                  self.assertEqual(node.kwargs["input"], ph_a)
                  if validate_same_kwargs:
                      for key, value in kwargs.items():
                          self.assertEqual(node.kwargs[key], value)
                  acc_op_node = node
              elif node.op == "output":
                  if acc_op is None:
                      # If we expect no new acc_op after graph building
                      # and found we have only output in traced graph
                      continue
                  self.assertEqual(acc_op_node, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          ref_outputs = m(a)
          outputs = traced(a)
          traced_again = acc_tracer.trace(traced, [a])
          outputs_again = traced_again(a)
          if isinstance(ref_outputs, torch.Tensor):
              ref_outputs = [ref_outputs]
              outputs = [outputs]
              outputs_again = [outputs_again]
  
          for ref_output, output, output_again in zip(
              ref_outputs, outputs, outputs_again
          ):
              if enable_allclose:
                  torch.testing.assert_close(
                      torch.nan_to_num(ref_output), torch.nan_to_num(output)
                  )
                  torch.testing.assert_close(
                      torch.nan_to_num(ref_output), torch.nan_to_num(output_again)
                  )
              else:
                  self.assertTrue(
                      torch.equal(torch.nan_to_num(ref_output), torch.nan_to_num(output))
                  )
                  self.assertTrue(
                      torch.equal(
                          torch.nan_to_num(ref_output), torch.nan_to_num(output_again)
                      )
                  )
  
      def test_sum(self):
          self._make_acc_op_function_test(acc_ops.sum, torch.sum)
          self._make_acc_op_function_test(acc_ops.sum, torch.sum, dim=(1,), keepdim=True)
  
      def test_prod(self):
          self._make_acc_op_function_test(acc_ops.prod, torch.prod)
          self._make_acc_op_function_test(acc_ops.prod, torch.prod, dim=1, keepdim=True)
  
      def test_mean(self):
          self._make_acc_op_function_test(acc_ops.mean, torch.mean)
          self._make_acc_op_function_test(
              acc_ops.mean, torch.mean, dim=(1,), keepdim=True
          )
  
      def test_pad(self):
          self._make_acc_op_function_test(
              acc_ops.pad, torch.nn.functional.pad, pad=(2, 0)
          )
  
      def test_max(self):
          def torch_max(x, *args, **kwargs):
              return x.max(*args, **kwargs)
  
          self._make_acc_op_function_test(acc_ops.max_full_reduce, torch_max)
          self._make_acc_op_function_test(
              acc_ops.max_dim_reduce, torch_max, dim=1, keepdim=True
          )
          self._make_acc_op_function_test(
              acc_ops.max_dim_reduce, torch_max, input_shape=(1, 4), dim=1, keepdim=True
          )
          self._make_acc_op_function_test(
              acc_ops.max_dim_reduce, torch_max, input_shape=(3, 4, 3), dim=2
          )
  
      @parameterized.expand(
          [
              param("max_maximum", orig_op=torch.max, expected_op=acc_ops.maximum),
              param(
                  "maximum_maximum", orig_op=torch.maximum, expected_op=acc_ops.maximum
              ),
              param("min_minimum", orig_op=torch.min, expected_op=acc_ops.minimum),
              param(
                  "minimum_minimum", orig_op=torch.minimum, expected_op=acc_ops.minimum
              ),
          ]
      )
      def test_maximum_minimum(self, _: str, orig_op, expected_op):
          class TestModule(torch.nn.Module):
              def __init__(self, orig_op):
                  super().__init__()
                  self.orig_op = orig_op
  
              def forward(self, input: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
                  return self.orig_op(input, other)
  
          m = TestModule(orig_op)
          input, other = torch.randn(2, 2), torch.randn(2, 2)
          traced = acc_tracer.trace(m, [input, other])
  
          ph_in = ph_oth = mxm = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  if str(node.target) == "other":
                      ph_oth = node
                  else:
                      self.assertTrue(str(node.target) == "input")
                      ph_in = node
              elif node.op == "call_function":
                  if node.target == expected_op:
                      self.assertEqual(node.kwargs["input"], ph_in)
                      self.assertEqual(node.kwargs["other"], ph_oth)
                      mxm = node
              elif node.op == "output":
                  self.assertEqual(mxm, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input, other), traced(input, other)))
  
      def test_conv(self):
          """
          Test that a conv is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.Conv2d(8, 7, 3, stride=2)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.randn(3, 8, 10, 10)
          traced = acc_tracer.trace(m, [input])
  
          ph = weight_attr = bias_attr = conv = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv.bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.conv2d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  self.assertEqual(node.kwargs["stride"], (2, 2))
                  self.assertEqual(node.kwargs["padding"], (0, 0))
                  self.assertEqual(node.kwargs["dilation"], (1, 1))
                  self.assertEqual(node.kwargs["groups"], 1)
                  conv = node
              elif node.op == "output":
                  self.assertEqual(conv, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_quantized_conv2d(self):
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.quantized.Conv2d(3, 3, 1)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.quantize_per_tensor(
              torch.randn(1, 3, 1, 1), scale=0.01, zero_point=3, dtype=torch.quint8
          )
          traced = acc_tracer.trace(m, [input])
          _LOGGER.info(traced.graph)
          ph = weight_attr = bias_attr = conv = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv_weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv_bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.quantized_conv2d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  conv = node
              elif node.op == "output":
                  self.assertEqual(conv, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_quantized_convrelu2d(self):
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.intrinsic.quantized.ConvReLU2d(3, 3, 1)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.quantize_per_tensor(
              torch.randn(1, 3, 1, 1), scale=0.01, zero_point=3, dtype=torch.quint8
          )
          traced = acc_tracer.trace(m, [input])
          ph = weight_attr = bias_attr = conv = relu = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv_weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv_bias":
                  bias_attr = node
              elif node.op == "call_function" and node.target == acc_ops.quantized_conv2d:
                  self.assertEqual(node.target, acc_ops.quantized_conv2d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  conv = node
              elif node.op == "call_function" and node.target == acc_ops.relu:
                  self.assertEqual(node.target, acc_ops.relu)
                  self.assertEqual(node.kwargs["input"], conv)
                  relu = node
              elif node.op == "output":
                  self.assertEqual(relu, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_conv1d(self):
          """
          Test that a conv is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.Conv1d(8, 7, 3, stride=2)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.randn(3, 8, 8)
          traced = acc_tracer.trace(m, [input])
  
          ph = weight_attr = bias_attr = conv = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv.bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.conv1d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  self.assertEqual(node.kwargs["stride"], (2,))
                  self.assertEqual(node.kwargs["padding"], (0,))
                  self.assertEqual(node.kwargs["dilation"], (1,))
                  self.assertEqual(node.kwargs["groups"], 1)
                  conv = node
              elif node.op == "output":
                  self.assertEqual(conv, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_conv3d(self):
          """
          Test that a conv is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.Conv3d(8, 7, 3, stride=2)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.randn(3, 8, 8, 10, 10)
          traced = acc_tracer.trace(m, [input])
  
          ph = weight_attr = bias_attr = conv = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv.bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.conv3d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  self.assertEqual(node.kwargs["stride"], (2, 2, 2))
                  self.assertEqual(node.kwargs["padding"], (0, 0, 0))
                  self.assertEqual(node.kwargs["dilation"], (1, 1, 1))
                  self.assertEqual(node.kwargs["groups"], 1)
                  conv = node
              elif node.op == "output":
                  self.assertEqual(conv, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_conv_transpose2d(self):
          """
          Test that a conv_transpose2d is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.ConvTranspose2d(8, 7, 3, stride=2)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.randn(3, 8, 10, 10)
          traced = acc_tracer.trace(m, [input])
  
          ph = weight_attr = bias_attr = conv = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv.bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.conv_transpose2d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  self.assertEqual(node.kwargs["stride"], (2, 2))
                  self.assertEqual(node.kwargs["padding"], (0, 0))
                  self.assertEqual(node.kwargs["output_padding"], (0, 0))
                  self.assertEqual(node.kwargs["groups"], 1)
                  self.assertEqual(node.kwargs["dilation"], (1, 1))
                  conv = node
              elif node.op == "output":
                  self.assertEqual(conv, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_conv_transpose3d(self):
          """
          Test that a conv_transpose3d is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.conv = nn.ConvTranspose3d(8, 7, 3, stride=2)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.conv(a)
  
          m = TestModule()
          input = torch.randn(3, 8, 8, 10, 10)
          traced = acc_tracer.trace(m, [input])
  
          ph = weight_attr = bias_attr = conv = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "conv.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "conv.bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.conv_transpose3d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  self.assertEqual(node.kwargs["stride"], (2, 2, 2))
                  self.assertEqual(node.kwargs["padding"], (0, 0, 0))
                  self.assertEqual(node.kwargs["output_padding"], (0, 0, 0))
                  self.assertEqual(node.kwargs["dilation"], (1, 1, 1))
                  self.assertEqual(node.kwargs["groups"], 1)
                  conv = node
              elif node.op == "output":
                  self.assertEqual(conv, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_embedding_bag(self):
          """
          Test that an embedding_bag is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.eb = nn.EmbeddingBag(10, 3, mode="sum", include_last_offset=True)
  
              def forward(self, inp: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
                  return self.eb(inp, offsets)
  
          m = TestModule()
          inp = torch.LongTensor([1, 2, 4, 5, 4, 3, 2, 9])
          offsets = torch.LongTensor([0, 4])
          traced = acc_tracer.trace(m, [inp, offsets])
  
          inp_node = offsets_node = weight_attr = eb_node = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  if str(node.target) == "inp":
                      inp_node = node
                  elif str(node.target) == "offsets":
                      offsets_node = node
                  else:
                      self.fail(f"Unexpected placeholder {node.target}.")
                  continue
              elif node.op == "get_attr" and node.target == "eb.weight":
                  weight_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.embedding_bag)
                  # Note: Normalization called from acc_tracer means we use all kwargs.
                  self.assertEqual(node.kwargs["input"], inp_node)
                  self.assertEqual(node.kwargs["offsets"], offsets_node)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["mode"], "sum")
                  self.assertEqual(node.kwargs["include_last_offset"], True)
                  # The rest of these were unspecified, so verify they fell back
                  # to their respective default values thanks to normalization.
                  self.assertEqual(node.kwargs["max_norm"], None)
                  self.assertEqual(node.kwargs["norm_type"], 2.0)
                  self.assertEqual(node.kwargs["scale_grad_by_freq"], False)
                  self.assertEqual(node.kwargs["sparse"], False)
                  self.assertEqual(node.kwargs["per_sample_weights"], None)
                  eb_node = node
              elif node.op == "output":
                  self.assertEqual(eb_node, node.args[0])
  
          self.assertTrue(torch.equal(m(inp, offsets), traced(inp, offsets)))
  
      def test_embedding_bag_byte_and_4bit_rowwise_offsets(self):
          """
          Test that 4 bit quantized embedding_bag is traced as expected.
          """
  
          class TestModule(nn.Module):
              def __init__(
                  self,
                  op,
                  q_weights,
                  per_index_weights,
              ):
                  super().__init__()
                  self.emb = op
                  self.q_weights = q_weights
                  self.per_index_weights = per_index_weights
  
              def forward(
                  self,
                  indices,
                  offsets,
              ):
                  return self.emb(
                      self.q_weights,
                      indices,
                      offsets,
                      mode=0,
                      per_sample_weights=self.per_index_weights,
                      include_last_offset=True,
                  )
  
          def run_embedding_bag_test(is_4bit, use_weights):
              # generate random indices, offsets, and weights.
              num_embeddings = 16
              embedding_dim = 32
              num_lengths = 10
  
              weights = torch.from_numpy(
                  (np.random.random_sample((num_embeddings, embedding_dim)) + 1).astype(
                      np.float32
                  )
              )
              q_weights = (
                  torch.ops.quantized.embedding_bag_4bit_prepack(weights)
                  if is_4bit
                  else torch.ops.quantized.embedding_bag_byte_prepack(weights)
              )
              np_lengths = np.random.randint(0, num_lengths, size=10).astype(np.int32)
  
              num_lengths = np.sum(np_lengths)
              indices = torch.from_numpy(
                  np.random.randint(low=0, high=num_embeddings, size=num_lengths)
              ).int()
  
              lengths = torch.from_numpy(np_lengths)
              offsets = torch.cat([torch.zeros([1]), torch.cumsum(lengths, 0)]).int()
  
              weights = torch.randint(low=0, high=4, size=indices.size())
              per_sample_weights = weights.to(torch.float32)
  
              indices = indices.to(torch.int32)
              offsets = offsets.to(torch.int32)
              inputs = [
                  indices,
                  offsets,
              ]
  
              op = (
                  torch.ops.quantized.embedding_bag_4bit_rowwise_offsets
                  if is_4bit
                  else torch.ops.quantized.embedding_bag_byte_rowwise_offsets
              )
  
              m = TestModule(
                  op,
                  q_weights,
                  per_sample_weights,
              )
  
              traced = acc_tracer.trace(m, inputs)
              _LOGGER.info(traced.graph)
  
              expected_target = (
                  acc_ops.embedding_bag_4bit_rowwise_offsets
                  if is_4bit
                  else acc_ops.embedding_bag_byte_rowwise_offsets
              )
  
              for node in traced.graph.nodes:
                  if node.op == "placeholder":
                      if str(node.target) == "indices":
                          inp_node = node
                      elif str(node.target) == "offsets":
                          offsets_node = node
                      else:
                          self.fail(f"Unexpected placeholder {node.target}.")
                      continue
                  elif node.op == "get_attr" and node.target == "q_weights":
                      weight_attr = node
                  elif node.op == "call_function":
                      self.assertEqual(node.target, expected_target)
                      # Note: Normalization called from acc_tracer means we use all kwargs.
                      self.assertEqual(node.kwargs["indices"], inp_node)
                      self.assertEqual(node.kwargs["offsets"], offsets_node)
                      self.assertEqual(node.kwargs["weight"], weight_attr)
                      self.assertEqual(node.kwargs["mode"], 0)
                      self.assertEqual(node.kwargs["include_last_offset"], True)
                      # The rest of these were unspecified, so verify they fell back
                      # to their respective default values thanks to normalization.
                      eb_node = node
                  elif node.op == "output":
                      self.assertEqual(eb_node, node.args[0])
              self.assertTrue(torch.equal(m(indices, offsets), traced(indices, offsets)))
  
          # test 8-bit
          run_embedding_bag_test(is_4bit=False, use_weights=True)
          # test 4-bit
          run_embedding_bag_test(is_4bit=True, use_weights=True)
  
      def test_quantized_batch_norm2d(self):
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.bn = nn.quantized.BatchNorm2d(3)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.bn(a)
  
          m = TestModule()
          m.eval()
          input = torch.quantize_per_tensor(
              torch.randn(1, 3, 1, 1), scale=0.01, zero_point=3, dtype=torch.quint8
          )
          traced = acc_tracer.trace(m, [input])
          ph = weight_attr = bias_attr = bn_mean = bn_var = bn = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "bn.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "bn.bias":
                  bias_attr = node
              elif node.op == "get_attr" and node.target == "bn.running_mean":
                  bn_mean = node
              elif node.op == "get_attr" and node.target == "bn.running_var":
                  bn_var = node
              elif node.op == "get_attr" and node.target == "bn.scale":
                  bn_scale = node
              elif node.op == "get_attr" and node.target == "bn.zero_point":
                  bn_zero_point = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.quantized_batch_norm2d)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  self.assertEqual(node.kwargs["running_mean"], bn_mean)
                  self.assertEqual(node.kwargs["running_var"], bn_var)
                  self.assertEqual(node.kwargs["acc_out_ty"][6]["scale"], bn_scale)
                  self.assertEqual(
                      node.kwargs["acc_out_ty"][6]["zero_point"], bn_zero_point
                  )
                  bn = node
              elif node.op == "output":
                  self.assertEqual(bn, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      def test_linear(self):
          """
          Test that a linear is traced as expected, i.e. to the functional level and with
          kwarg normalization. Also verify that symbolic shape inference worked as part of
          the acc_tracer.
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.linear = nn.Linear(3, 5, bias=True)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.linear(a)
  
          m = TestModule()
          test_input = torch.randn(1, 3)
          traced = acc_tracer.trace(m, [test_input])
          ph = weight_attr = bias_attr = linear = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "linear.weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "linear.bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.linear)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  linear = node
              elif node.op == "output":
                  self.assertEqual(linear, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
          self.assertTrue(torch.equal(m(test_input), traced(test_input)))
  
      def test_quantized_linear(self):
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
                  self.linear = nn.quantized.Linear(3, 5)
  
              def forward(self, a: torch.Tensor) -> torch.Tensor:
                  return self.linear(a)
  
          m = TestModule()
          input = torch.quantize_per_tensor(
              torch.randn(2, 3), scale=0.01, zero_point=3, dtype=torch.quint8
          )
          traced = acc_tracer.trace(m, [input])
          ph = weight_attr = bias_attr = linear = None
          for node in traced.graph.nodes:
              if node.op == "placeholder":
                  self.assertEqual(str(node.target), "a")
                  ph = node
              elif node.op == "get_attr" and node.target == "linear_weight":
                  weight_attr = node
              elif node.op == "get_attr" and node.target == "linear_bias":
                  bias_attr = node
              elif node.op == "call_function":
                  self.assertEqual(node.target, acc_ops.quantized_linear)
                  self.assertEqual(node.kwargs["input"], ph)
                  self.assertEqual(node.kwargs["weight"], weight_attr)
                  self.assertEqual(node.kwargs["bias"], bias_attr)
                  linear = node
              elif node.op == "output":
                  self.assertEqual(linear, node.args[0])
              else:
                  self.fail(f"Unexpected node: {node.format_node()}")
  
          self.assertTrue(torch.equal(m(input), traced(input)))
  
      @parameterized.expand(
          [
              param("remove_exceptions_false", remove_exceptions=False),
              param("remove_exceptions_true", remove_exceptions=True),
          ]
      )
      def test_batch_norm(self, _, remove_exceptions):
          """
          Test that a batch norm is traced as expected, i.e. to the functional level
          and with kwarg normalization. Note that we also expect to see a
          ConditionalExceptionWrapper in the graph that the AST rewriter converted
          from `if x: raise y`.
  
          """
  
          class TestModule(nn.Module):
              def __init__(self):
                  super().__init__()
  