class_docstr: Base TestCase to be used for all conv-bn[-relu] fusion patterns.
class_docstr_tok: ["Base","TestCase","to","be","used","for","all","conv-bn","[","-relu","]","fusion","patterns","."]
docstr: Test whether `source_fn_stack` is preserved after QAT fusion.
docstr_tok: ["Test","whether","`","source_fn_stack","`","is","preserved","after","QAT","fusion","."]
code: def test_qat_preserve_source_fn_stack(self):
      """
      Test whether `source_fn_stack` is preserved after QAT fusion.
      """
  
      class M(torch.nn.Module):
          def __init__(self, conv_class, bn_class, backbone):
              super().__init__()
              self.conv = conv_class(5, 3, 3)
              self.bn = bn_class(3)
              self.relu = torch.nn.ReLU()
              self.backbone = backbone
  
          def forward(self, x):
              x = self.conv(x)
              x = self.bn(x)
              x = self.relu(x)
              x = self.backbone(x)
              return x
  
      assert self.dim in [1, 2]
      if self.dim == 1:
          example_inputs = (torch.randn(1, 5, 10),)
      else:
          example_inputs = (torch.randn(1, 5, 10, 10),)
  
      # QAT prepare + convert
      backbone = self._get_conv_bn_model(has_relu=True)
      m = M(self.conv_class, self.bn_class, backbone)
      quantizer = XNNPACKQuantizer()
      quantizer.set_global(
          get_symmetric_quantization_config(is_qat=True)
      )
      m = capture_pre_autograd_graph(m, example_inputs)
      m = prepare_qat_pt2e(m, quantizer)
      m(*example_inputs)
      m = convert_pt2e(m)
  
      # Extract the conv and relu nodes (bn was folded into conv)
      first_conv, first_relu, second_conv, second_relu = (
          None,
          None,
          None,
          None,
      )
      for n in m.graph.nodes:
          if n.target == torch.ops.aten.relu.default:
              if first_relu is None:
                  assert first_conv is None, "bad test setup"
                  first_relu = n
                  first_conv = n.args[0]
              else:
                  assert second_conv is None, "bad test setup"
                  second_relu = n
                  second_conv = n.args[0]
  
      # Extract the conv weight and bias nodes
      def get_conv_weight_and_bias(conv_node: torch.fx.Node):
          weight_dq_node = conv_node.args[1]
          weight_q_node = weight_dq_node.args[0]
          weight_node = weight_q_node.args[0]
          bias_node = conv_node.args[2]
          assert isinstance(weight_node, torch.fx.Node)
          assert isinstance(bias_node, torch.fx.Node)
          return (weight_node, bias_node)
  
      first_conv_weight, first_conv_bias = get_conv_weight_and_bias(
          first_conv
      )
      second_conv_weight, second_conv_bias = get_conv_weight_and_bias(
          second_conv
      )
  
      # Assert that each set of conv, conv weight, and conv bias are in the same partition
      def get_source_fn(node: torch.fx.Node):
          # E.g. [('l__self___backbone1_conv', <class 'torch.nn.modules.conv.Conv2d'>)]
          return node.meta["source_fn_stack"][0][0]
  
      self.assertEqual(
          get_source_fn(first_conv), get_source_fn(first_conv_weight)
      )
      self.assertEqual(
          get_source_fn(first_conv), get_source_fn(first_conv_bias)
      )
      self.assertEqual(
          get_source_fn(second_conv), get_source_fn(second_conv_weight)
      )
      self.assertEqual(
          get_source_fn(second_conv), get_source_fn(second_conv_bias)
      )
  
      # Assert that different sets of convs and relus have different partitions
      self.assertNotEqual(
          get_source_fn(first_conv), get_source_fn(first_relu)
      )
      self.assertNotEqual(
          get_source_fn(first_conv), get_source_fn(second_conv)
      )
      self.assertNotEqual(
          get_source_fn(second_conv), get_source_fn(second_relu)
      )
      self.assertNotEqual(
          get_source_fn(first_relu), get_source_fn(second_relu)
      )
  
      # Assert that "backbone" exists only in the second set of conv and relu's partition
      self.assertTrue("backbone" not in get_source_fn(first_conv))
      self.assertTrue("backbone" not in get_source_fn(first_relu))
      self.assertTrue("backbone" in get_source_fn(second_conv))
      self.assertTrue("backbone" in get_source_fn(second_relu))
  
imports: import copy
  import operator
  import unittest
  from typing import Any, Optional, Tuple, Type
  import torch
  from torch._export import capture_pre_autograd_graph
  from torch.ao.quantization import default_fake_quant, FusedMovingAvgObsFakeQuantize, MovingAverageMinMaxObserver, MovingAveragePerChannelMinMaxObserver, QConfigMapping
  from torch.ao.quantization.backend_config import get_qnnpack_backend_config
  from torch.ao.quantization.qconfig import default_per_channel_symmetric_qnnpack_qat_qconfig, default_symmetric_qnnpack_qat_qconfig
  from torch.ao.quantization.quantize_fx import prepare_qat_fx
  from torch.ao.quantization.quantize_pt2e import _convert_to_reference_decomposed_fx, convert_pt2e, prepare_pt2e, prepare_qat_pt2e
  from torch.ao.quantization.quantizer import DerivedQuantizationSpec, QuantizationAnnotation, QuantizationSpec, Quantizer
  from torch.ao.quantization.quantizer.xnnpack_quantizer import get_symmetric_quantization_config, XNNPACKQuantizer
  from torch.testing._internal.common_cuda import TEST_CUDA
  from torch.testing._internal.common_quantization import NodeSpec, QuantizationTestCase, skip_if_no_torchvision, skipIfNoQNNPACK
  from torch.testing._internal.common_quantized import override_quantized_engine
label_desc: Use image node_id 13 for calling the TestQuantizePT2EQAT_ConvBn_Base obj's underlying member method code with example usage: obj.test_qat_preserve_source_fn_stack() and returns: weight_node,  bias_node, node, x