docstr: Tests that graphs containing branches are prepared correctly.
  Specifically, equalization observers should not be inserted in front of
  branches in which both initial layers in the branches plan to be
  quantized.
docstr_tok: ["Tests","that","graphs","containing","branches","are","prepared","correctly",".","Specifically",",","equalization","observers","should","not","be","inserted","in","front","of","branches","in","which","both","initial","layers","in","the","branches","plan","to","be","quantized","."]
code: def test_input_weight_equalization_branching(self):
      """Tests that graphs containing branches are prepared correctly.
      Specifically, equalization observers should not be inserted in front of
      branches in which both initial layers in the branches plan to be
      quantized.
      """
  
      # Tests that we do not add an equalization observer due to both initial
      # nodes in the branch containing layers that need to be equalized.
      # Note that this should print out 2 warning messages for not being able
      # to equalize layers linear1 and linear1 because it is part of a branch
      class TestBranchingWithoutEqualizationModel(nn.Module):
          def __init__(self) -> None:
              super().__init__()
              self.linear1 = nn.Linear(5, 5)
              self.linear2 = nn.Linear(5, 5)
  
          def forward(self, x):
              y = self.linear1(x)
              z = self.linear2(x)
              return torch.add(y, z)
  
      no_eq_branching_node_occurrence = {
          ns.call_module(_InputEqualizationObserver): 0,
          ns.call_module(MinMaxObserver): 3,
      }
  
      m = TestBranchingWithoutEqualizationModel().eval()
      example_inputs = (torch.rand(1, 5),)
      prepared = prepare_fx(
          m,
          specific_qconfig_dict,
          example_inputs=example_inputs,
          _equalization_config=default_equalization_qconfig_dict,
      )
      self.checkGraphModuleNodes(
          prepared,
          expected_node_occurrence=no_eq_branching_node_occurrence,
      )
  
      # Tests that we will add an equalization observer because there is only
      # one initial node in the branch that needs to be equalized
      class TestBranchingWithEqualizationModel(nn.Module):
          def __init__(self) -> None:
              super().__init__()
              self.linear1 = nn.Linear(5, 5)
  
          def forward(self, x):
              y = self.linear1(x)
              z = torch.add(x, 5)
              return torch.add(y, z)
  
      eq_branching_node_occurrence = {
          ns.call_module(_InputEqualizationObserver): 1,
          ns.call_module(MinMaxObserver): 2,
      }
  
      m = TestBranchingWithEqualizationModel().eval()
      example_inputs = (torch.randn(1, 5),)
      prepared = prepare_fx(
          m,
          specific_qconfig_dict,
          example_inputs=example_inputs,
          _equalization_config=default_equalization_qconfig_dict,
      )
      self.checkGraphModuleNodes(
          prepared,
          expected_node_occurrence=eq_branching_node_occurrence,
      )
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  import torch.ao.nn.intrinsic.quantized
  import torch.ao.nn.quantized
  from torch.ao.quantization import default_qconfig
  from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver
  from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx
  from torch.ao.quantization.fx._equalize import _InputEqualizationObserver, _WeightEqualizationObserver, calculate_equalization_scale, default_equalization_qconfig, _convert_equalization_ref, get_layer_sqnr_dict, get_equalization_qconfig_dict
  from torch.testing._internal.common_quantization import NodeSpec, QuantizationTestCase, SingleLayerLinearModel, TwoLayerLinearModel, LinearAddModel, SingleLayerFunctionalLinearModel, TwoLayerFunctionalLinearModel, FunctionalLinearAddModel, ConvModel, TwoLayerConvModel, SingleLayerFunctionalConvModel, TwoLayerFunctionalConvModel, skipIfNoFBGEMM, LinearReluModel, LinearReluLinearModel, LinearReluAddModel, FunctionalLinearReluModel, FunctionalLinearReluLinearModel, ConvReluModel, ConvReluConvModel, ConvReluAddModel, FunctionalConvReluModel, FunctionalConvReluConvModel
  import copy
  import numpy
  from hypothesis import given
  from hypothesis import strategies
label_desc: Use image node_id 4 for calling the TestEqualizeFx obj's underlying member method code with example usage: obj.test_input_weight_equalization_branching() and returns: torch, torch