class_docstr: Prepares the model for QAT by inserting weight and activation quantizers as
  specified in qconfig_mapping.
  
  Implements additional graph passes on a prepared module returned by prepare_qat_fx.
class_docstr_tok: ["Prepares","the","model","for","QAT","by","inserting","weight","and","activation","quantizers","as","specified","in","qconfig_mapping",".","Implements","additional","graph","passes","on","a","prepared","module","returned","by","prepare_qat_fx","."]
docstr: Some layers like ReLU can be quantized with affine qscheme even when we want
  to use symmetric quantization (zero point = 0). This is because these layers
  always have a non-negative output. And thus, an affine activation post process layer attached
  after layers like these will always observe zero point as 0. This can possibly help us
  reduce quantization error because of the larger number of quantization levels available.
  (Symmetric quantization will force the output of these layers to use [0, 127] as the
  output range, but with affine quantization, we can use [0, 255]).
  
  prepare_qat_fx requires all modules being fused together to have the same QConfig.
  Thus, if we have a Conv followed by a ReLU and we want to set ReLU to have affine qscheme,
  we would have to set Conv to use affine qscheme as well. But this isn't correct because a stand alone
  Conv layer somewhere else in the network will also use affine qscheme which is undesirable
  we want to fix zero point to 0.
  
  Hence, we add this pass which replaces all occurrences of activation post process after
   ``always_affine_layers`` with an affine version.
docstr_tok: ["Some","layers","like","ReLU","can","be","quantized","with","affine","qscheme","even","when","we","want","to","use","symmetric","quantization","(","zero","point","=","0",")",".","This","is","because","these","layers","always","have","a","non-negative","output",".","And","thus",",","an","affine","activation","post","process","layer","attached","after","layers","like","these","will","always","observe","zero","point","as","0",".","This","can","possibly","help","us","reduce","quantization","error","because","of","the","larger","number","of","quantization","levels","available",".","(","Symmetric","quantization","will","force","the","output","of","these","layers","to","use","[","0",",","127","]","as","the","output","range",",","but","with","affine","quantization",",","we","can","use","[","0",",","255","]",")",".","prepare_qat_fx","requires","all","modules","being","fused","together","to","have","the","same","QConfig",".","Thus",",","if","we","have","a","Conv","followed","by","a","ReLU","and","we","want","to","set","ReLU","to","have","affine","qscheme",",","we","would","have","to","set","Conv","to","use","affine","qscheme","as","well",".","But","this","is","n't","correct","because","a","stand","alone","Conv","layer","somewhere","else","in","the","network","will","also","use","affine","qscheme","which","is","undesirable","we","want","to","fix","zero","point","to","0",".","Hence",",","we","add","this","pass","which","replaces","all","occurrences","of","activation","post","process","after","``","always_affine_layers","``","with","an","affine","version","."]
code: def _mark_always_affine_layers_for_replacement(
      self, model: _fx.GraphModule
  ):
      """
      Some layers like ReLU can be quantized with affine qscheme even when we want
      to use symmetric quantization (zero point = 0). This is because these layers
      always have a non-negative output. And thus, an affine activation post process layer attached
      after layers like these will always observe zero point as 0. This can possibly help us
      reduce quantization error because of the larger number of quantization levels available.
      (Symmetric quantization will force the output of these layers to use [0, 127] as the
      output range, but with affine quantization, we can use [0, 255]).
  
      prepare_qat_fx requires all modules being fused together to have the same QConfig.
      Thus, if we have a Conv followed by a ReLU and we want to set ReLU to have affine qscheme,
      we would have to set Conv to use affine qscheme as well. But this isn't correct because a stand alone
      Conv layer somewhere else in the network will also use affine qscheme which is undesirable
      we want to fix zero point to 0.
  
      Hence, we add this pass which replaces all occurrences of activation post process after
       ``always_affine_layers`` with an affine version.
      """
      # Note: For all these ops, whether or not we can use affine qscheme for them depends only on
      # the op itself or one preceding op.
      # Note: graph.nodes traverses the nodes in topological order
      for node in model.graph.nodes:
          if node.op == "call_module":
              layer = _find_target(model, node.target)
              if type(layer) in _always_affine_layers:
                  self._mark_act_post_process_for_replacement(
                      node, model
                  )
              elif isinstance(
                  layer,
                  (_qat.ConvAct2d, _qat.ConvBnAct2d, _qat.LinearAct),
              ):
                  if type(layer.act) in _always_affine_layers:
                      self._mark_act_post_process_for_replacement(
                          node, model
                      )
              # layers which only scale the output can also use affine qcheme
              elif isinstance(layer, tuple(_scale_only_layers)):
                  arg_mod = _find_module(model, node.args[0])
                  if (
                      _is_activation_post_process(arg_mod)
                      and node.args[0] in self._modules_to_replace
                  ):
                      self._mark_act_post_process_for_replacement(
                          node, model
                      )
          elif node.op == "call_function":
              combine_op_type = _combine_op_type(node)
              if combine_op_type is not None:
                  if combine_op_type == _CombinationOpType.AddReLU:
                      self._mark_act_post_process_for_replacement(
                          node, model
                      )
              elif node.target in _always_affine_layers:
                  self._mark_act_post_process_for_replacement(
                      node, model
                  )
  
imports: from collections import defaultdict
  from typing import Any
  from typing import Optional
  import torch
  import torch.ao.quantization
  import torch.fx
  import torch.nn
  import torch.nn.intrinsic
  import torch.nn.intrinsic.qat
  from torch.ao.quantization.backend_config import BackendConfig
  from torch.ao.quantization.fx.custom_config import PrepareCustomConfig
  from torch.quantization.quantize_fx import prepare_qat_fx
  import coremltools.optimize.torch.quantization.modules.qat_modules
  from coremltools.optimize.torch._utils.torch_utils import get_parent_child_name
  from coremltools.optimize.torch.quantization._backend_config import _fixed_qparams_modules
  from coremltools.optimize.torch.quantization._utils import CombinationOpType
  from coremltools.optimize.torch.quantization._utils import combine_op_type
  from coremltools.optimize.torch.quantization._utils import find_module
  from coremltools.optimize.torch.quantization._utils import find_target
  from coremltools.optimize.torch.quantization._utils import get_share_qparams_ops
  from coremltools.optimize.torch.quantization._utils import group_activation_quantization_modules_by_id
  from coremltools.optimize.torch.quantization._utils import is_activation_post_process
  from coremltools.optimize.torch.quantization._utils import is_quantized
  from coremltools.optimize.torch.quantization.quantization_config import QuantizationScheme
label_desc: Use image node_id 6 for calling the QATConfigurationHandler obj's underlying member method code with example usage: obj._mark_always_affine_layers_for_replacement(model) without return types