text: """
   Copyright (c) 2019-2020 Intel Corporation
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
        http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
  """
  
  # pylint:disable=too-many-lines
  import functools
  from collections import OrderedDict, namedtuple
  from copy import deepcopy
  from enum import Enum
  from pathlib import Path
  from typing import List, Dict, Tuple, Optional, Callable
  
  import networkx as nx
  import numpy as np
  import operator
  import shutil
  import torch
  from nncf.layers import NNCFEmbedding
  from torch import nn
  
  from nncf.algo_selector import COMPRESSION_ALGORITHMS
  from nncf.compression_method_api import CompressionAlgorithmBuilder, CompressionAlgorithmController, CompressionLevel
  from nncf.debug import is_debug, DebugInterface, CallCountTracker
  from nncf.dynamic_graph.context import OperatorInput, TracingContext, Scope
  from nncf.dynamic_graph.function_input_quantization import FUNCTIONS_TO_QUANTIZE
  from nncf.dynamic_graph.graph import NNCFNode, InputAgnosticOperationExecutionContext
  from nncf.dynamic_graph.graph import NNCFNodeExpression as N, NNCFGraph
  from nncf.dynamic_graph.patch_pytorch import get_arg_positions_to_quantize, MODEL_INPUT_OP_NAME
  from nncf.dynamic_graph.transform_graph import is_nncf_module
  from nncf.hw_config import HWConfig
  from nncf.initialization import DataLoaderRangeInitializeRunner
  from nncf.module_operations import UpdateWeight, UpdateInputs
  from nncf.nncf_logger import logger as nncf_logger
  from nncf.nncf_network import NNCFNetwork, CompressionModuleType, InsertionCommand, OperationPriority, \
      InsertionPoint, InsertionType, InsertionPointGraph, InsertionPointGraphNodeType
  from nncf.quantization.hw_precision_constraints import HWPrecisionConstraints
  from nncf.quantization.init_precision import PrecisionInitializerFactory
  from nncf.quantization.layers import QUANTIZATION_MODULES, QuantizationMode, QuantizerConfig, BaseQuantizer, \
      QuantizerExportMode, QuantizersSwitcher
  from nncf.quantization.metrics import NetworkQuantizationShareMetric, MemoryCostMetric, ShareEdgesQuantizedDataPath
  from nncf.quantization.quantizer_id import WeightQuantizerId, NonWeightQuantizerId, InputQuantizerId, \
      FunctionQuantizerId
  from nncf.quantization.quantizer_propagation import QuantizerPropagationSolver, QuantizerPropagationStateGraph, \
      QuantizersBetweenQuantizableLayers, QuantizerInsertionInfo
  from nncf.quantization.schedulers import QUANTIZATION_SCHEDULERS
  from nncf.structures import QuantizationPrecisionInitArgs, QuantizationRangeInitArgs
  from nncf.utils import get_all_modules_by_type, in_scope_list, is_main_process, should_consider_scope
  from nncf.utils import get_state_dict_names_with_modules
  
  
  class QuantizerSetupType(Enum):
      PATTERN_BASED = "pattern_based"
      PROPAGATION_BASED = "propagation_based"
  
      @staticmethod
      def from_str(quantizer_setup_type: str) -> 'QuantizerSetupType':
          if quantizer_setup_type == QuantizerSetupType.PATTERN_BASED.value:
              return QuantizerSetupType.PATTERN_BASED
          if quantizer_setup_type == QuantizerSetupType.PROPAGATION_BASED.value:
              return QuantizerSetupType.PROPAGATION_BASED
          raise RuntimeError("Unknown quantizer setup type. Please select 'pattern_based' or 'propagation_based'.")
  
  
  class QuantizationConstraints:
      REF_QCONF_OBJ = QuantizerConfig()
  
      def __init__(self, **kwargs):
          """Use attribute names of QuantizerConfig as arguments
          to set up constraints.
          E.g. QuantizationConstraint(bits=8, per_channel=True) will set up
          a constraint that corresponds to all 8-bit per-channel quantizers, either
          symmetric or asymmetric, either signed or unsigned."""
  
          for attr_name in kwargs:
              if not hasattr(QuantizationConstraints.REF_QCONF_OBJ, attr_name):
                  raise RuntimeError("Invalid constraint - QuantizerConfig has no attribute '{}'".format(attr_name))
          self.qconf_attr_vs_constraint_dict = kwargs
  
      def apply_constraints_to(self, qconfig: QuantizerConfig) -> QuantizerConfig:
          for attr_name, constraint in self.qconf_attr_vs_constraint_dict.items():
              if constraint is not None:
                  setattr(qconfig, attr_name, constraint)
          return qconfig
  
      def is_config_compatible(self, qconfig: QuantizerConfig) -> bool:
          is_compatible = True
          for attr_name, constraint in self.qconf_attr_vs_constraint_dict.items():
              if constraint is not None:
                  qconf_attr_value = getattr(qconfig, attr_name)
                  if qconf_attr_value != constraint:
                      is_compatible = False
          return is_compatible
  
  
  class QuantizerGroup(Enum):
      ACTIVATIONS = "activations"
      WEIGHTS = "weights"
  
  
  PotentialQuantizedModule = namedtuple('PotentialQuantizedModule', 'module module_scope qconfig_list')
  
  
  class NonWeightQuantizerInfo:
      def __init__(self, quantizer_module_ref: BaseQuantizer,
                   affected_ia_op_exec_contexts: List[InputAgnosticOperationExecutionContext],
                   quantizers_between_quantizable_layers: QuantizersBetweenQuantizableLayers = None):
          self.quantizer_module_ref = quantizer_module_ref
          self.affected_ia_op_exec_contexts = affected_ia_op_exec_contexts
          self.quantizers_between_quantizable_layers = quantizers_between_quantizable_layers
  
  
  @COMPRESSION_ALGORITHMS.register('quantization')
  class QuantizationBuilder(CompressionAlgorithmBuilder):
      DEFAULT_QUANTIZER_CONFIG = QuantizerConfig(bits=8,
                                                 mode=QuantizationMode.SYMMETRIC,
                                                 signedness_to_force=None,
                                                 per_channel=False)
  
      def __init__(self, config, should_init: bool = True):
          super().__init__(config, should_init)
  
          self.quantize_inputs = self.config.get('quantize_inputs', True)
          self.quantize_outputs = self.config.get('quantize_outputs', False)
          self.disable_function_quantization_hooks = self.config.get('disable_function_quantization_hooks', False)
          self._debug_interface = QuantizationDebugInterface() if is_debug() else None
  
          self._quantized_weight_modules_registry = OrderedDict()
          self._quantized_inputs_modules_registry = OrderedDict()
          self._weight_quantizers = OrderedDict()  # Quantizers applied via UpdateWeights
          self._non_weight_quantizers = OrderedDict()  # All the other quantizers
          self._processed_function_quantizers = set()
          self._processed_input_agnostic_op_exec_contexts = set()
  
          self.global_quantizer_contraints = {}  # type: Dict[QuantizerGroup, QuantizationConstraints]
          self._ignored_scopes_per_group = {}  # type: Dict[QuantizerGroup, List[str]]
          self._target_scopes_per_group = {}  # type: Dict[QuantizerGroup, List[str]]
  
          for quantizer_group in QuantizerGroup:
              self._parse_group_params(self.config, quantizer_group)
  
          self.quantizer_setup_type = self.config.get('quantizer_setup_type')
          self.quantizable_subgraph_patterns = self.config.get('quantizable_subgraph_patterns', None)
          self.hw_config = None
          hw_config_type = self.config.get("hw_config_type")
          is_hw_config_enabled = hw_config_type is not None
          self._hw_precision_constraints = HWPrecisionConstraints(is_hw_config_enabled)
          if is_hw_config_enabled:
              hw_config_path = HWConfig.get_path_to_hw_config(hw_config_type)
              self.hw_config = HWConfig.from_json(hw_config_path)
  
      def _parse_group_params(self, quant_config: 'NNCFConfig', quantizer_group: QuantizerGroup):
          group_name = quantizer_group.value
          params_dict = quant_config.get(group_name, {})
          self.global_quantizer_contraints[quantizer_group] = QuantizationConstraints(
              bits=params_dict.get('bits'),
              mode=params_dict.get('mode'),
              signedness_to_force=params_dict.get('signed'),
              per_channel=params_dict.get('per_channel')
          )
          self._ignored_scopes_per_group[quantizer_group] = params_dict.get('ignored_scopes')
          self._target_scopes_per_group[quantizer_group] = params_dict.get('target_scopes')
  
      def apply_to(self, target_model: NNCFNetwork) -> NNCFNetwork:
          insertion_commands = self._quantize_weights(target_model) + self._quantize_activations(target_model)
          if self.quantize_inputs and self.quantizer_setup_type is not QuantizerSetupType.PROPAGATION_BASED:
              insertion_commands += self._quantize_inputs(target_model, insertion_commands)
  
          quantization_types = [class_type.__name__ for class_type in QUANTIZATION_MODULES.registry_dict.values()]
  
          # At this point the NNCF module quantization modules are not in the target_model yet,
          # therefore it is extended with the corresponding registries tracked during weights/inputs quantizations
          self._all_quantizations = get_state_dict_names_with_modules(target_model, quantization_types)
          self._all_quantizations.update(self._quantized_weight_modules_registry)
          self._all_quantizations.update(self._quantized_inputs_modules_registry)
  
          for command in insertion_commands:
              target_model.register_insertion_command(command)
  
          target_model.register_algorithm(self)
  
          if self._debug_interface is not None:
              target_model.debug_interface.add_interface(self._debug_interface)
          return target_model
  
      def build_controller(self, target_model: NNCFNetwork) -> CompressionAlgorithmController:
          return QuantizationController(target_model,
                                        self.config,
                                        self.should_init,
                                        self._debug_interface,
                                        self._quantized_weight_modules_registry,
                                        self._quantized_inputs_modules_registry,
                                        self._weight_quantizers,
                                        self._non_weight_quantizers,
                                        self._hw_precision_constraints)
  
      def __get_default_qconfig(self, constraints: QuantizationConstraints = None):
          qconfig = deepcopy(self.DEFAULT_QUANTIZER_CONFIG)
          if constraints is not None:
              qconfig = constraints.apply_constraints_to(qconfig)
          return qconfig
  
      def __get_scoped_quantizer_config(self, target_model: NNCFNetwork,
                                        parent_module_scope_str: str, is_weights=False,
                                        input_shape=None) -> QuantizerConfig:
          group = QuantizerGroup.WEIGHTS if is_weights else QuantizerGroup.ACTIVATIONS
          qconfig = self.__get_default_qconfig(constraints=self.global_quantizer_contraints[group])
          qconfig.is_weights = is_weights
          scope_overrides = self.config.get("scope_overrides", {})
          for overridden_scope in scope_overrides.keys():
              if in_scope_list(parent_module_scope_str, overridden_scope):
                  config_overrides = scope_overrides[overridden_scope]
                  if config_overrides.get("bits") is not None:
                      qconfig.bits = config_overrides["bits"]
                  if config_overrides.get("mode") is not None:
                      qconfig.mode = config_overrides["mode"]
                  if config_overrides.get("per_channel") is not None:
                      qconfig.per_channel = config_overrides["per_channel"]
                  if config_overrides.get("signed") is not None:
                      qconfig.signedness_to_force = config_overrides["signed"]
          if qconfig.per_channel:
              if is_weights:
                  module = target_model.get_module_by_scope(Scope.from_str(parent_module_scope_str))
                  qconfig.input_shape = module.weight.shape
              elif input_shape is not None:
                  qconfig.input_shape = input_shape
              else:
                  raise RuntimeError("Unable to use per channel quantization for module {} activations -"
                                     " input shape is unknown".format(parent_module_scope_str))
          return qconfig
  
      def __create_quantize_module(self, qconfig: QuantizerConfig):
          quantizer_cls = QUANTIZATION_MODULES.get(qconfig.mode)
          return quantizer_cls(qconfig)
  
      def _make_quantizable_subgraph_pattern(self):
          full_pattern = self._make_default_quantizable_subgraph_pattern()
          if self.quantizable_subgraph_patterns is not None:
              for pattern in self.quantizable_subgraph_patterns:
                  if not isinstance(pattern, str):
                      custom_pattern = functools.reduce(operator.add,
                                                        [N(node) for node in pattern])
                  else:
                      custom_pattern = N(pattern)
                  full_pattern = full_pattern | custom_pattern
          return full_pattern
  
      def _check_if_ip_graph_nodes_point_to_single_module(self, ip_graph_node_list: List[dict]):
          """Does not access actual modules - only uses the InputAgnosticOperationExecutionContext info."""
          ia_op_exec_contexts_list = []  # type: List[InputAgnosticOperationExecutionContext]
          for ip_graph_op_node in ip_graph_node_list:
              nncf_node = ip_graph_op_node[InsertionPointGraph.REGULAR_NODE_REF_NODE_ATTR]
              ia_op_exec_context = nncf_node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].input_agnostic
              ia_op_exec_contexts_list.append(ia_op_exec_context)
  
          contexts_correspond_to_single_module = True
          first_op_context = ia_op_exec_contexts_list[0]
          for other_op_context in ia_op_exec_contexts_list:
              if other_op_context.scope_in_model != first_op_context.scope_in_model or \
                      other_op_context.operator_name != first_op_context.operator_name:
                  contexts_correspond_to_single_module = False
                  break
  
          if not contexts_correspond_to_single_module:
              raise RuntimeError("NNCF module has more than 1 associated graph operation node corresponding"
                                 "to different module hierarchy locations - cannot make sure that weight "
                                 "quantization will be correct")
  
      def get_potential_quantized_modules(self, target_model: NNCFNetwork) -> List[PotentialQuantizedModule]:
          modules = target_model.get_nncf_modules()
          insertion_point_graph = target_model.get_insertion_point_graph()
          quantized_modules_with_potential_qconfig = []
          default_qconfig_list = [self.__get_default_qconfig(
              constraints=self.global_quantizer_contraints[QuantizerGroup.WEIGHTS])]
          if self.hw_config is not None:
              meta_vs_qconfig_map = self.hw_config.get_metatype_vs_quantizer_configs_map(for_weights=True)
          for module_scope, module in modules.items():
              if not self._should_consider_scope_for_group(str(module_scope), QuantizerGroup.WEIGHTS):
                  nncf_logger.info("Ignored adding Weight quantizer in scope: {}".format(module_scope))
                  continue
              if self.hw_config is None:
                  qconfig = self.__get_scoped_quantizer_config(target_model, str(module_scope), is_weights=True)
                  qconfig_list = [qconfig]
              else:
                  associated_ops = insertion_point_graph.get_op_nodes_in_scope(module_scope)
                  if not associated_ops:
                      raise RuntimeError(
                          "Could not find a patched operation corresponding to NNCF module scope {}".format(
                              str(module_scope)))
  
                  if len(associated_ops) > 1:
                      self._check_if_ip_graph_nodes_point_to_single_module(associated_ops)
                  graph_operation = associated_ops[0]
                  metatype = graph_operation[InsertionPointGraph.OPERATOR_METATYPE_NODE_ATTR]
                  qconfig_list = meta_vs_qconfig_map[metatype]
                  if HWConfig.is_wildcard_quantization(qconfig_list):  # Empty list = wildcard quantization
                      qconfig_list = default_qconfig_list
  
              if qconfig_list is not None:
                  quantized_modules_with_potential_qconfig.append(PotentialQuantizedModule(module, module_scope,
                                                                                           qconfig_list))
          return quantized_modules_with_potential_qconfig
  
      def _quantize_weights(self, target_model: NNCFNetwork) -> List[InsertionCommand]:
          device = next(target_model.parameters()).device
          insertion_commands = []
          quantized_modules_with_potential_qconfig = \
              self.get_potential_quantized_modules(target_model)
          for module, module_scope, qconfig_list in quantized_modules_with_potential_qconfig:
              self._quantized_weight_modules_registry[str(module_scope)] = module
              nncf_logger.info("Adding signed Weight quantizer in scope: {}".format(module_scope))
  
              if self.hw_config is not None:
                  try:
                      qconfig_overrides = None
                      if in_scope_list(str(module_scope), self.config.get("scope_overrides", {})):
                          qconfig_overrides = self.__get_scoped_quantizer_config(target_model,
                                                                                 str(module_scope),
                                                                                 is_weights=True)
                      qconfig = self._select_final_qconfig(qconfig_list,
                                                           self.global_quantizer_contraints[QuantizerGroup.WEIGHTS],
                                                           qconfig_overrides)
                  except RuntimeError:
                      err_msg = "Quantization parameter constraints specified in NNCF config are incompatible with HW "
                      err_msg += "capabilities as specified in HW config type '{}'. ".format(self.hw_config.target_device)
                      err_msg += "First conflicting quantizer location: {}".format(str(module_scope))
                      raise RuntimeError(err_msg)
              else:
                  assert len(
                      qconfig_list) == 1, "Non-HW config scenarios should produce single quantizer configs for each " \
                                          "weight module!"
                  qconfig = qconfig_list[0]
  
              quantizer_id = WeightQuantizerId(module_scope)
              self._hw_precision_constraints.add(quantizer_id, qconfig_list)
              qconfig.input_shape = module.weight.shape
              quantizer = self.__create_quantize_module(qconfig)
              op = UpdateWeight(quantizer).to(device)
              # TODO: separate insertion point semantic for weights and activations
              insertion_commands.append(InsertionCommand(InsertionPoint(
                  InputAgnosticOperationExecutionContext('', module_scope, 0),
                  InsertionType.NNCF_MODULE_PRE_OP), op, OperationPriority.QUANTIZATION_PRIORITY))
              self._weight_quantizers[quantizer_id] = quantizer
          return insertion_commands
  
      class ActivationQuantizationHook:
          """Cannot simply register the quantizer module as a callable hook, since we need to call
          a thread-local version of the quantizer module during base module execution."""
  
          def __init__(self, context: TracingContext, quantizer_storage_key: str,
                       debug_interface: 'QuantizationDebugInterface' = None):
              self.compressed_context = context
              self.quantizer_storage_key = quantizer_storage_key
              self.debug_interface = debug_interface
  
          def __call__(self, *args, **kwargs):
              if self.debug_interface is not None:
                  self.debug_interface.register_activation_quantize_call(str(self.quantizer_storage_key))
              replica = self.compressed_context.base_module_thread_local_replica
              return replica.activation_quantizers[self.quantizer_storage_key](*args, **kwargs)
  
      def _quantize_activations(self, target_model: NNCFNetwork) -> List[InsertionCommand]:
          target_model.register_compression_module_type(CompressionModuleType.ACTIVATION_QUANTIZER)
  
          if self.quantizer_setup_type == QuantizerSetupType.PATTERN_BASED:
              insertion_commands = self._quantize_post_pattern_activations(target_model)
          elif self.quantizer_setup_type == QuantizerSetupType.PROPAGATION_BASED:
              insertion_point_graph = target_model.get_insertion_point_graph()
              if self._debug_interface:
                  self._debug_interface.visualize_insertion_point_graph(insertion_point_graph)
              prop_graph_solver = QuantizerPropagationSolver(ignored_scopes=self.ignored_scopes,
                                                             debug_interface=self._debug_interface,
                                                             hw_config=self.hw_config,
                                                             default_qconfig_list=[self.__get_default_qconfig(
                                                                 constraints=self.global_quantizer_contraints[
                                                                     QuantizerGroup.ACTIVATIONS])],
                                                             input_infos=target_model.get_input_infos())
              merged_ip_graph = insertion_point_graph.get_ip_graph_with_merged_hw_optimized_operations(self.hw_config)
              insertion_data = prop_graph_solver.run_on_ip_graph(merged_ip_graph)
              insertion_commands = []
  
              original_nncf_graph = target_model.get_original_graph()
              for insertion_info, quantizer_config_list in insertion_data.items():
                  ia_op_exec_context = insertion_info.op_exec_context.input_agnostic
                  operator_scope_str = str(ia_op_exec_context)
                  if not self.quantize_inputs and ia_op_exec_context.operator_name == MODEL_INPUT_OP_NAME:
                      continue
  
                  # Tailored for post-hook quantization and first output quantization only
                  quantizer_input_shape = original_nncf_graph.get_output_shapes_for_ia_op_exec_context(
                      ia_op_exec_context)[0]
                  try:
                      qconfig_overrides = None
                      if in_scope_list(operator_scope_str, self.config.get("scope_overrides", {})):
                          qconfig_overrides = self.__get_scoped_quantizer_config(target_model,
                                                                                 operator_scope_str,
                                                                                 is_weights=True)
                      quantizer_config = self._select_final_qconfig(quantizer_config_list,
                                                                    self.global_quantizer_contraints[
                                                                        QuantizerGroup.ACTIVATIONS],
                                                                    qconfig_overrides)
                  except RuntimeError:
                      err_msg = "Quantization parameter constraints specified in NNCF config are incompatible with HW "
                      err_msg += "capabilities as specified in HW config type '{}'. ".format(self.hw_config.target_device)
                      err_msg += "First conflicting quantizer location: "
                      err_msg += str(ia_op_exec_context)
                      raise RuntimeError(err_msg)
  
                  quantizer_config.input_shape = quantizer_input_shape
                  quantizer_id = NonWeightQuantizerId(ia_op_exec_context)
                  self._hw_precision_constraints.add(quantizer_id, quantizer_config_list)
                  insertion_commands += self._add_single_activation_quantizer(target_model, insertion_info,
                                                                              quantizer_config)
          else:
              raise RuntimeError("Invalid quantizer setup type!")
  
          if not self.disable_function_quantization_hooks:
              insertion_commands += self._quantize_free_function_inputs(target_model)
          return insertion_commands
  
      def _select_final_qconfig(self, quantizer_config_list: List[QuantizerConfig],
                                constraints: QuantizationConstraints, qconfig_overrides=None) -> QuantizerConfig:
          assert quantizer_config_list is not None
  
          if self.hw_config is None and qconfig_overrides is not None:
              return qconfig_overrides
  
          if HWConfig.is_wildcard_quantization(quantizer_config_list):
              # Set a default, most basic quantization config in case wildcard propagating quantizer did
              # not merge to align with other tools
              return self.__get_default_qconfig()
  
          constrained_quantizer_config_list = list(filter(
              constraints.is_config_compatible,
              quantizer_config_list
          ))
  
  
          if qconfig_overrides is not None:
              constrained_quantizer_config_list =\
                   [qconfig for qconfig in constrained_quantizer_config_list if qconfig_overrides == qconfig]
  
          # TODO: Make the logic more flexible when the flag "warning as error" is implemented.
          # It means that the qconfig from overrides must be selected as final config
          # even if it is not valid in hw-config.
          if not constrained_quantizer_config_list:
              raise RuntimeError()
  
          # Quantizer config list entries should arrive in the same order as they are listed
          # in the HW config, where they are sorted by descending order of priority
          return constrained_quantizer_config_list[0]
  
      def _quantize_post_pattern_activations(self, target_model: NNCFNetwork) -> List[InsertionCommand]:
          pattern = self._make_quantizable_subgraph_pattern()
          target_insertion_infos = target_model.get_post_pattern_insertion_points(pattern)
          insertion_commands = []
  
          act_config = self.config.get('activations', {})
          if 'linked_quantizer_scopes' in act_config:
              linked_scopes_groups_list = act_config['linked_quantizer_scopes']
              target_insertion_infos = self.coalesce_insertion_infos(target_insertion_infos, linked_scopes_groups_list)
  
          for insertion_info in target_insertion_infos:
              ia_op_exec_context = insertion_info.op_exec_context.input_agnostic
              operator_scope_str = str(ia_op_exec_context)
  
              if not self.quantize_outputs and insertion_info.is_output:
                  nncf_logger.info("Ignored adding Activation Quantize "
                                   "in scope (output scope, quantize_outputs=False): {}".format(operator_scope_str))
                  continue
              if not self._should_consider_scope_for_group(operator_scope_str, QuantizerGroup.ACTIVATIONS):
                  nncf_logger.info("Ignored adding Activation quantizer in scope: {}".format(operator_scope_str))
                  continue
  
              qconfig = self.__get_scoped_quantizer_config(target_model, operator_scope_str,
                                                           is_weights=False,
                                                           input_shape=insertion_info.shape_to_operate_on)
              quantizer_insertion_info = QuantizerInsertionInfo.from_insertion_info(insertion_info)
              insertion_commands += self._add_single_activation_quantizer(target_model, quantizer_insertion_info, qconfig)
  
          # NOTE: Order of activations must be the same to correctly broadcast parameters (e.g. scales) in distributed
          # mode (see call of `_dist_broadcast_coalesced` in torch/nn/parallel/distributed.py for more details)
          # pylint: disable=protected-access
          target_model.sort_compression_modules(CompressionModuleType.ACTIVATION_QUANTIZER)
          return insertion_commands
  
      def _add_single_activation_quantizer(self, target_model: NNCFNetwork,
                                           insertion_info: QuantizerInsertionInfo,
                                           quantizer_config: QuantizerConfig) -> List[InsertionCommand]:
          """Will return one or more insertion commands - depending on whether insertion_info specifies
          a single input agnostic operation execution context or there are linked contexts along with it."""
          ia_op_exec_context = insertion_info.op_exec_context.input_agnostic
          operator_scope_str = str(ia_op_exec_context)
          device = next(target_model.parameters()).device
          quantizer = self.__create_quantize_module(quantizer_config).to(device)
          affected_ia_op_exec_contexts = [ia_op_exec_context] + [x.input_agnostic for x in
                                                                 insertion_info.linked_op_exec_contexts]
  
          # linked_op_exec_contexts will determine quantization points in graph that have to share
          # a quantization module, e.g. for scales unification
          serialized_context_list = [str(x) for x in affected_ia_op_exec_contexts]
          quantizer_storage_key = ";".join([str(x) for x in serialized_context_list])
  
          assert quantizer_storage_key not in target_model.get_compression_modules_by_type(
              CompressionModuleType.ACTIVATION_QUANTIZER)
  
          target_model.add_compression_module(quantizer_storage_key, quantizer,
                                              CompressionModuleType.ACTIVATION_QUANTIZER)
  
          quantizer_id = NonWeightQuantizerId(ia_op_exec_context)
  
          if len(affected_ia_op_exec_contexts) > 1:
              nncf_logger.info(
                  "Processing linked activation quantizer group:\n {}\n".format("\n".join(serialized_context_list)))
  
          self._non_weight_quantizers[quantizer_id] = NonWeightQuantizerInfo(
              quantizer, affected_ia_op_exec_contexts, insertion_info.quantizers_between_quantizable_layers)
  
          insertion_commands = []
          for curr_ia_op_exec_context in affected_ia_op_exec_contexts:
              if curr_ia_op_exec_context in self._processed_input_agnostic_op_exec_contexts:
                  raise RuntimeError(
                      "Ambiguous call to {fn} with call order {co} in current scope. "
                      "Cannot insert quantization hooks "
                      "automatically!".format(fn=ia_op_exec_context.operator_name, co=ia_op_exec_context.call_order)
                  )
              self._processed_input_agnostic_op_exec_contexts.add(curr_ia_op_exec_context)
  
              nncf_logger.info("Adding {} Activation Quantize in scope: {}".format(
                  "signed" if quantizer.signed else
                  "unsigned", operator_scope_str
              ))
  
              # Hooks will be identical for each affected ia_op_exec_context - will call one and the
              # same quantizer
              hook = self.ActivationQuantizationHook(target_model.get_tracing_context(),
                                                     quantizer_storage_key,
                                                     self._debug_interface)
  
              insertion_commands.append(InsertionCommand(InsertionPoint(curr_ia_op_exec_context,
                                                                        InsertionType.OPERATOR_POST_HOOK),
                                                         hook,
                                                         OperationPriority.QUANTIZATION_PRIORITY))
          return insertion_commands
  
      def _quantize_inputs(self, target_model: NNCFNetwork,
                           prev_weight_and_activation_quantizer_insertion_commands: List[InsertionCommand]) -> \
          List[InsertionCommand]:
          device = next(target_model.parameters()).device
          graph_roots = target_model.get_original_graph().get_input_nodes()
  
          # Have to handle the situation when the input node of the network is an NNCF module -
          # to quantize inputs in this case we will have to use UpdateInputs module pre-op,
  
          # Traverse down starting from graph roots to search for the first node which belongs to a NNCF module
          # and has no UpdateInputs pre-op
  
          def traverse_function(node: NNCFNode, output) -> Tuple[bool, List[NNCFNode]]:
              module = target_model.get_module_by_scope(node.op_exec_context.scope_in_model)
              if is_nncf_module(module):
                  if isinstance(module, NNCFEmbedding):
                      # Embeddings have integer input and their quantization is rather controlled
                      # by their weights.
                      return True, output
                  current_node_scope = node.op_exec_context.scope_in_model
                  module_op_insertion_commands = []
                  for comm in prev_weight_and_activation_quantizer_insertion_commands:
                      if current_node_scope in comm.insertion_point.ia_op_exec_context.scope_in_model:
                          module_op_insertion_commands.append(comm)
                  pre_op_insertion_commands = filter(
                      lambda comm: comm.insertion_point.insertion_type == InsertionType.NNCF_MODULE_PRE_OP,
                      module_op_insertion_commands)
                  update_inputs_count = sum(1 for comm in pre_op_insertion_commands if isinstance(comm.fn, UpdateInputs))
                  if update_inputs_count == 0:
                      output.append(node)
                      return True, output
              else:
                  current_node_ia_op_exec_context = node.op_exec_context.input_agnostic
                  op_hook_insertion_commands = []
                  for comm in prev_weight_and_activation_quantizer_insertion_commands:
                      if current_node_ia_op_exec_context == comm.insertion_point.ia_op_exec_context:
                          op_hook_insertion_commands.append(comm)
                  if op_hook_insertion_commands:
                      return True, output
  
              return False, output
  
          nncf_module_input_nodes = set()
          for node in graph_roots:
              scope_str = str(node.op_exec_context.scope_in_model)
              if self._should_consider_scope_for_group(scope_str, QuantizerGroup.ACTIVATIONS):
                  nncf_module_input_nodes.update(
                      target_model.get_original_graph().traverse_graph(node, traverse_function))
  
          insertion_commands = []
          nncf_scope_module_dict = target_model.get_nncf_modules()
          for module_input_node in nncf_module_input_nodes:
              op_scope = module_input_node.op_exec_context.input_agnostic.scope_in_model
              module = None
              scope = None
              for nncf_scope, nncf_module in nncf_scope_module_dict.items():
                  if op_scope in nncf_scope:
                      module = nncf_module
                      scope = nncf_scope
                      break
  
              self._quantized_inputs_modules_registry[str(scope)] = module
  
              # Only use the shape of the 0-th input info specified in config. TODO: fix this
              input_shape = target_model.input_infos[0].shape if target_model.input_infos is not None else None
              qconfig = self.__get_scoped_quantizer_config(target_model, str(scope), is_weights=False,
                                                           input_shape=input_shape)
              quantizer = self.__create_quantize_module(qconfig)
  
              nncf_logger.info("Adding {} NNCF module input quantizer in scope: {}".format(
                  "signed" if quantizer.signed else "unsigned", str(scope)
              ))
  
              # TODO: separate insertion point semantic for weights and activations
              insertion_commands.append(
                  InsertionCommand(InsertionPoint(InputAgnosticOperationExecutionContext("", scope, 0),
                                                  InsertionType.NNCF_MODULE_PRE_OP),
                                   UpdateInputs(quantizer).to(device),
                                   OperationPriority.QUANTIZATION_PRIORITY))
              ia_op_exec_context = module_input_node.op_exec_context.input_agnostic
              quantizer_id = InputQuantizerId(ia_op_exec_context)
              self._hw_precision_constraints.add(quantizer_id, [qconfig])
              self._non_weight_quantizers[quantizer_id] = NonWeightQuantizerInfo(
                  quantizer, [ia_op_exec_context], QuantizersBetweenQuantizableLayers())
  
          return insertion_commands
  
      def _should_consider_scope_for_group(self, scope_str: str, group: QuantizerGroup) -> bool:
          if self.target_scopes is not None or self._target_scopes_per_group[group] is not None:
              if in_scope_list(scope_str, self.target_scopes):
                  return True
              if in_scope_list(scope_str, self._target_scopes_per_group[group]):
                  return True
  
              return False
  
          if in_scope_list(scope_str, self.ignored_scopes):
              return False
          if in_scope_list(scope_str, self._ignored_scopes_per_group[group]):
              return False
  
          return True
  
      class FunctionQuantizationPreHook:
          """Cannot simply register the quantizer module as a callable hook, since we need to call
          a thread-local version of the quantizer module during base module execution."""
  
          def __init__(self, context: TracingContext, func_in_quant_info: FunctionQuantizerId,
                       debug_interface: 'QuantizationDebugInterface' = None):
              self.compressed_context = context
              self.func_in_quant_info = func_in_quant_info
              self.debug_interface = debug_interface
  
          def __call__(self, op_inputs: OperatorInput):
              quantizer_dict_key = str(self.func_in_quant_info)
              if self.debug_interface is not None:
                  self.debug_interface.register_function_quantizer_call(quantizer_dict_key)
              replica = self.compressed_context.base_module_thread_local_replica
              idx = self.func_in_quant_info.input_arg_idx
              op_inputs.op_args[idx] = replica.function_quantizers[quantizer_dict_key](op_inputs.op_args[idx])
              return op_inputs
  
      def _quantize_free_function_inputs(self, target_model: NNCFNetwork) -> List[InsertionCommand]:
          device = next(target_model.parameters()).device
  
          if not FUNCTIONS_TO_QUANTIZE:
              return []
          pattern = N(FUNCTIONS_TO_QUANTIZE[0].name)
          for i in range(1, len(FUNCTIONS_TO_QUANTIZE)):
              pattern |= N(FUNCTIONS_TO_QUANTIZE[i].name)
  
          target_insertion_infos = target_model.get_post_pattern_insertion_points(pattern,
                                                                                  omit_nodes_in_nncf_modules=True)
          insertion_commands = []
  
          target_model.register_compression_module_type(CompressionModuleType.FUNCTION_QUANTIZER)
          for insertion_info in target_insertion_infos:
              ia_op_exec_context = insertion_info.op_exec_context.input_agnostic
              scope_str = str(ia_op_exec_context.scope_in_model)
  
              if not self._should_consider_scope_for_group(scope_str, QuantizerGroup.ACTIVATIONS):
                  nncf_logger.info("Ignored adding function input quantizer in scope: {}".format(scope_str))
                  continue
  
              function_arg_positions_to_quantize = get_arg_positions_to_quantize(ia_op_exec_context.operator_name)
              assert function_arg_positions_to_quantize is not None, "Function with inputs to be quantized has " \
                                                                     "no info struct registered in " \
                                                                     "QUANTIZED_INPUT_FUNCTIONS!"
  
              for input_arg_idx in function_arg_positions_to_quantize:
                  ip_arg_quant_key = FunctionQuantizerId(ia_op_exec_context, input_arg_idx)
  
                  if ip_arg_quant_key in self._processed_function_quantizers:
                      raise RuntimeError(
                          "Ambiguous call to {fn} with call order {co} and argname {arg} in current scope. "
                          "Cannot insert quantization hooks "
                          "automatically!".format(fn=ia_op_exec_context.operator_name,
                                                  co=ia_op_exec_context.call_order,
                                                  arg=input_arg_idx)
                      )
  
                  self._processed_function_quantizers.add(ip_arg_quant_key)
  
                  ip_arg_quant_name = str(ip_arg_quant_key)
                  assert ip_arg_quant_name not in target_model.get_compression_modules_by_type(
                      CompressionModuleType.FUNCTION_QUANTIZER)
                  input_shape = insertion_info.op_exec_context.tensor_metas[0].shape
  
                  qconfig = self.__get_scoped_quantizer_config(target_model, scope_str,
                                                               is_weights=False,
                                                               input_shape=input_shape)
                  quantizer_module = self.__create_quantize_module(qconfig).to(device)
                  target_model.add_compression_module(ip_arg_quant_name, quantizer_module,
                                                      CompressionModuleType.FUNCTION_QUANTIZER)
  
                  nncf_logger.info("Adding {} Function Quantize: {}".format(
                      "signed" if quantizer_module.signed else
                      "unsigned", ip_arg_quant_name))
  
                  hook = self.FunctionQuantizationPreHook(target_model.get_tracing_context(),
                                                          ip_arg_quant_key,
                                                          self._debug_interface)
                  insertion_commands.append(InsertionCommand(InsertionPoint(ia_op_exec_context,
                                                                            InsertionType.OPERATOR_PRE_HOOK),
                                                             hook,
                                                             OperationPriority.QUANTIZATION_PRIORITY))
                  self._hw_precision_constraints.add(ip_arg_quant_key, [qconfig])
                  self._non_weight_quantizers[ip_arg_quant_key] = NonWeightQuantizerInfo(quantizer_module,
                                                                                         [ia_op_exec_context])
          # NOTE: Order of input quantizers must be the same to correctly broadcast parameters (e.g. scales) in
          # distributed mode (see call of `_dist_broadcast_coalesced` in torch/nn/parallel/distributed.py for more
          # details) pylint: disable=protected-access
          target_model.sort_compression_modules(CompressionModuleType.FUNCTION_QUANTIZER)
          return insertion_commands
  
      @staticmethod
      def _make_default_quantizable_subgraph_pattern():
          import nncf.dynamic_graph.patterns as p
          pattern = p.LINEAR_OPS | p.ARITHMETIC | p.ANY_BN_ACT_COMBO | \
                    p.LINEAR_OPS + p.ANY_BN_ACT_COMBO | p.ARITHMETIC + p.ANY_BN_ACT_COMBO | p.SINGLE_OPS | p.MATMUL
          return pattern
  
      @staticmethod
      def coalesce_insertion_infos(target_insertion_infos: List[QuantizerInsertionInfo],
                                   linked_scopes_groups_list: List[List[str]]) -> List[QuantizerInsertionInfo]:
          """Accepts a list of InsertionInfos that each correspond only to one InputAgnosticOperationExecutionContext,
          and merges these according to linked_scope_groups_list so that some or all of the resulting InsertionInfo
          objects have non-empty linked_op_exec_contexts lists.
          Each entry in linked_scope_groups_list must be a valid string representation of a single
          InputAgnosticOperationExecutionContext object."""
          ia_op_exec_context_list = [x.op_exec_context.input_agnostic for x in target_insertion_infos]
          retval = []
          insertion_info_indices_vs_group_id = OrderedDict()
  
          for group_idx, group_list in enumerate(linked_scopes_groups_list):
              for group_member_scope_str in group_list:
                  ia_op_exec_context = InputAgnosticOperationExecutionContext.from_str(group_member_scope_str)
                  matching_indices = list(
                      filter(lambda x: ia_op_exec_context_list[x] == ia_op_exec_context,
                             range(len(ia_op_exec_context_list))))
                  if len(matching_indices) > 1:
                      raise RuntimeError(
                          "Linked activation quantizer entry {} specifies more than 1 activation quantizer:\n {}".format(
                              group_member_scope_str,
                              "\n".join([str(ia_op_exec_context_list[i]) for i in matching_indices])))
                  if len(matching_indices) == 0:
                      raise RuntimeError("No match for linked quantizer entry {} among activation quantizers!".format(
                          group_member_scope_str))
  
                  target_idx = matching_indices[0]
                  if target_idx in insertion_info_indices_vs_group_id:
                      raise RuntimeError(
                          "Linked activation quantizer groups {} and {} "
                          "overlap!".format(group_idx,
                                            insertion_info_indices_vs_group_id[target_idx])
                      )
                  insertion_info_indices_vs_group_id[target_idx] = group_idx
  
          for i in range(len(ia_op_exec_context_list)):
              if i not in insertion_info_indices_vs_group_id:
                  insertion_info_indices_vs_group_id[i] = None
  
          group_indices_list = [[] for _ in linked_scopes_groups_list]  # type: List[List[int]]
          for insertion_info_idx, group_idx in insertion_info_indices_vs_group_id.items():
              if group_idx is not None:
                  group_indices_list[group_idx].append(insertion_info_idx)
              else:
                  retval.append(target_insertion_infos[insertion_info_idx])
  
          for intra_group_indices in group_indices_list:
              main_info_idx = intra_group_indices[0]
              main_info = target_insertion_infos[main_info_idx]
              new_info = QuantizerInsertionInfo(main_info.op_exec_context,
                                                main_info.is_input,
                                                main_info.is_output,
                                                shape_to_operate_on=main_info.shape_to_operate_on)
              for linked_info_idx in intra_group_indices[1:]:
                  new_info.linked_op_exec_contexts.append(target_insertion_infos[linked_info_idx].op_exec_context)
              retval.append(new_info)
  
          return retval
  
  
  class QuantizationControllerBase(CompressionAlgorithmController):
      def enable_activation_quantization(self):
          raise NotImplementedError
  
      def enable_weight_quantization(self):
          raise NotImplementedError
  
      def disable_activation_quantization(self):
          raise NotImplementedError
  
      def disable_weight_quantization(self):
          raise NotImplementedError
  
      def init_range(self):
          raise NotImplementedError
  
  
  class QuantizationController(QuantizationControllerBase):
      def __init__(self, target_model: NNCFNetwork,
                   quantization_config: 'NNCFConfig',
                   should_init: bool,
                   debug_interface: 'QuantizationDebugInterface',
                   quantized_weight_modules_registry: Dict[Scope, torch.nn.Module],
                   quantized_inputs_modules_registry: Dict[Scope, torch.nn.Module],
                   weight_quantizers: Dict[WeightQuantizerId, torch.nn.Module],
                   non_weight_quantizers: Dict[NonWeightQuantizerId, NonWeightQuantizerInfo],
                   hw_precision_constraints: HWPrecisionConstraints,
                   collect_compression_metrics: bool = True):
          super().__init__(target_model)
          self.debug_interface = debug_interface
          self.quantization_config = quantization_config
          self._hw_precision_constraints = hw_precision_constraints
          self._collect_compression_metrics = collect_compression_metrics
  
          self.quantized_weight_modules_registry = quantized_weight_modules_registry
          self.quantized_inputs_modules_registry = quantized_inputs_modules_registry
          self.weight_quantizers = weight_quantizers
          self.non_weight_quantizers = non_weight_quantizers
          self.all_quantizations = OrderedDict()
          self.all_quantizations.update(self.weight_quantizers)
          self.all_quantizations.update({k: v.quantizer_module_ref for k, v in self.non_weight_quantizers.items()})
  
          should_export_to_onnx_qdq = quantization_config.get("export_to_onnx_standard_ops",
                                                              False)
          if should_export_to_onnx_qdq:
              export_mode = QuantizerExportMode.ONNX_QUANTIZE_DEQUANTIZE_PAIRS
          else:
              export_mode = QuantizerExportMode.FAKE_QUANTIZE
  
          for quantizer in self.all_quantizations.values():  # type: BaseQuantizer
              quantizer.set_export_mode(export_mode)
  
          if is_main_process() and should_init:
              self.initialize_quantizer_params()
  
          # Staged scheduler must be created after initialized to prevent extra logic with disabled quantizations
          params = quantization_config.get('params', None)
          self.is_staged_scheduler = bool(params)
          if self.is_staged_scheduler:
              scheduler_cls = QUANTIZATION_SCHEDULERS.get("staged")
              self._scheduler = scheduler_cls(self, params)
  
          if self._collect_compression_metrics:
              self.metric_store = {}
              quantizer_setup_type = self.quantization_config.get('quantizer_setup_type')
              # These metrics are collected here and are updated when the method .statistics() is called
              self.non_stable_metric_collectors = [NetworkQuantizationShareMetric(target_model, self.weight_quantizers, \
                                                                                  self.non_weight_quantizers,
                                                                                  quantizer_setup_type),
                                                   MemoryCostMetric(target_model, self.weight_quantizers,
                                                                    self.non_weight_quantizers)]
              # These metrics are collected once here and are not updated when the method .statistics() is called
              self.stable_metric_collectors = [ShareEdgesQuantizedDataPath(target_model)]
              self.update_metric_store(True)
  
      def prepare_for_export(self):
          for quantizer_id, quantizer in self.all_quantizations.items():
              if not quantizer.is_enabled_quantization():
                  nncf_logger.warning('Disabled quantization on export to ONNX: {}'.format(quantizer_id))
  
      def update_metric_store(self, do_all: bool = False):
          for collector in self.non_stable_metric_collectors:
              collector.collect()
              self.metric_store[collector.NAME_STR] = collector.get_metric_table()
          if do_all:
              for collector in self.stable_metric_collectors:
                  collector.collect()
                  self.metric_store[collector.NAME_STR] = collector.get_metric_table()
  
      def distributed(self):
          # NOTE: Order of quantization modules must be the same on GPUs to correctly broadcast num_bits
          sorted_quantizers = OrderedDict(sorted(self.all_quantizations.items(), key=lambda x: str(x[0])))
          for quantizer in sorted_quantizers.values():  # type: BaseQuantizer
              quantizer.broadcast_initialized_params()
  
      def _do_range_init(self, data_loader,
                         num_init_steps: int,
                         global_init_range_config: dict,
                         device: str):
          modules_to_init = OrderedDict()
          scope_overrides = self.quantization_config.get("scope_overrides", {})
  
          for class_type in QUANTIZATION_MODULES.registry_dict.values():
              quantization_type = class_type.__name__
              module_dict = get_all_modules_by_type(self._model, quantization_type)
              for scope, module in module_dict.items():
                  quantizer_group = "weights" if module.is_weights else "activations"
                  module_init_range_config = self._get_local_init_range_config(scope, scope_overrides,
                                                                               global_init_range_config, quantizer_group)
                  modules_to_init[str(scope)] = (module, module_init_range_config)
  
          # NOTE: Order of modules must be the same to correctly broadcast parameters (e.g. input_low
          # and input_range)
          modules_to_init = OrderedDict(sorted(modules_to_init.items()))
          self.modules_to_range_init = modules_to_init
          runner = DataLoaderRangeInitializeRunner(self._model, modules_to_init, device)
  
          quantizers = [module for module, config in modules_to_init.values()]
          quantizers_switcher = QuantizersSwitcher(quantizers)
          # bypass quantization to collect statistics from floating point model
          quantizers_switcher.disable_quantizers()
          runner.run(data_loader, num_init_steps)
          quantizers_switcher.enable_quantizers()
  
          self._model.rebuild_graph()
  
      def compression_level(self) -> CompressionLevel:
          if self.is_staged_scheduler:
              return self.scheduler.compression_level()
          return CompressionLevel.FULL
  
      def initialize_quantizer_params(self):
          """ For the quantization there are 2 types of initializations: range and precision
              BatchNorm statistics are updated for the quantized model as a final initialization step.
          """
          self.init_range()
          self.init_precision()
          self.run_batchnorm_adaptation(self.quantization_config)
  
      def init_precision(self):
          """
          Precision initialization happens based on an measure of layer sensitivity to perturbations. The measure is
          calculated by average Hessian trace estimation for each layer using Hutchinson algorithm.
          Parameters for the quantization algorithm:
              'data_loader' - provides an iterable over the given dataset, instance of 'torch.utils.data.DataLoader'
              'criterion' - loss function, instance of `torch.nn.modules.loss._Loss`,
          """
          init_config = self.quantization_config.get('initializer', {})
          init_precision_config = init_config.get('precision', None)
          if init_precision_config:
              precision_init_type = init_precision_config.get('type', 'manual')
              precision_init_args = None
              if precision_init_type != 'manual':
                  try:
                      precision_init_args = self.quantization_config.get_extra_struct(QuantizationPrecisionInitArgs)
                  except KeyError:
                      raise ValueError(
                          'Specified non-manual precision initialization in the NNCF config, '
                          'but the initializing data loader and loss criterion are not provided as an extra struct. '
                          'Refer to `NNCFConfig.register_extra_structs` and the `QuantizationPrecisionInitArgs` class')
  
              init_impl = PrecisionInitializerFactory.create(precision_init_type)
              initializer = init_impl(self, init_precision_config, precision_init_args)
              nncf_logger.info("Initialization of quantization precisions")
              initializer.apply_init()
  
      def init_range(self):
          """
          Calculates per-layer activation statistics on training dataset in order to choose proper output range for
          quantization.
          Parameters for quantization algorithm:
              'data_loader' - provides an iterable over the given dataset, instance of 'torch.utils.data.DataLoader'
          """
          init_config = self.quantization_config.get('initializer', {})
          init_range_config = init_config.get('range', {})
          if not init_range_config:
              try:
                  self.quantization_config.get_extra_struct(QuantizationRangeInitArgs)
                  has_range_init_args = True
              except KeyError:
                  has_range_init_args = False
  
              if has_range_init_args:
                  nncf_logger.warning("Enabling quantization range initialization with default parameters.")
                  num_init_steps = 1
              else:
                  nncf_logger.warning("Initializer section not specified for quantization algorithm in NNCF config and "
                                      "quantization init args not supplied - quantizer range initialization will not be "
                                      "done")
                  num_init_steps = 0
  
              init_range_config = {'num_init_steps': num_init_steps}
          if isinstance(init_range_config, dict):
              global_init_range_config = self.update_range_config_by_default(init_range_config)
              max_num_init_steps = global_init_range_config['num_init_steps']
          else:
              max_num_init_steps = 0
              global_init_range_config = []
              for sub_init_range_config in init_range_config:
                  global_init_range_config.append(self.update_range_config_by_default(sub_init_range_config))
                  max_num_init_steps = max(sub_init_range_config['num_init_steps'], max_num_init_steps)
  
          if max_num_init_steps > 0:
              try:
                  range_init_args = self.quantization_config.get_extra_struct(QuantizationRangeInitArgs)
              except KeyError:
                  raise ValueError(
                      'Should run range initialization as specified via config,'
                      'but the initializing data loader is not provided as an extra struct. '
                      'Refer to `NNCFConfig.register_extra_structs` and the `QuantizationRangeInitArgs` class')
              data_loader = range_init_args.data_loader
  
              self._do_range_init(data_loader, max_num_init_steps, global_init_range_config,
                                  range_init_args.device)
  
      def update_range_config_by_default(self, init_range_config: Dict):
          global_init_range_config = dict()
          global_init_range_config.update(init_range_config)
          if global_init_range_config.get("type") is None:
              global_init_range_config["type"] = "mean_min_max"
  
          if global_init_range_config.get("num_init_steps") is None:
              global_init_range_config["num_init_steps"] = 1
  
          num_init_steps = global_init_range_config.get('num_init_steps', 1)
          if num_init_steps < 0:
              raise AttributeError('Number of initialization steps must be >= 0')
          return global_init_range_config
  
      def get_weights_activation_quantizers_pairs(self) -> List[Tuple[List[BaseQuantizer], BaseQuantizer]]:
          """
          finds all neighbour weight and input activation quantizers that share the same module (e.g. conv or linear).
          Single activation quantizer can be in pair with multiple neighbour weight quantizers, e.g. like in SqueezeNet,
          when two Convolutions share the same input activation.
          :return: list of pairs - (list of weight quantizers, activation quantizer)
          """
          pairs = []
          qimr = OrderedDict(sorted(self.quantized_inputs_modules_registry.items()))
          for _, quantized_module in qimr.items():
              weight_quantizer = None
              activation_quantizer = None
              for ops in quantized_module.pre_ops.values():
                  if isinstance(ops, UpdateWeight):
                      weight_quantizer = ops.op
                  if isinstance(ops, UpdateInputs):
                      activation_quantizer = ops.op
              if weight_quantizer:
                  pairs.append(([weight_quantizer], activation_quantizer))
  
          nncf_network = self._model
          nncf_graph = nncf_network.get_original_graph()
          non_weight_quantizers = {key: quantizer_info.quantizer_module_ref for key, quantizer_info \
                                   in self.non_weight_quantizers.items() if not isinstance(key, InputQuantizerId)}
  
          def traverse_graph(curr_nx_node_key: str, weight_quantizers: List[nn.Module]) -> Optional[List[nn.Module]]:
              nx_node = nncf_graph.get_nx_node_by_key(curr_nx_node_key)
              module_scope = nx_node[NNCFGraph.OP_EXEC_CONTEXT_NODE_ATTR].scope_in_model
              module = nncf_network.get_module_by_scope(module_scope)
              if is_nncf_module(module):
                  if hasattr(module, 'pre_ops'):
                      for ops in module.pre_ops.values():
                          if isinstance(ops, UpdateWeight):
                              weight_quantizers.append(ops.op)
              else:
                  for succ_nx_node_key in nncf_graph.get_successors(curr_nx_node_key):
                      return traverse_graph(succ_nx_node_key, weight_quantizers)
              return weight_quantizers
  
          # pylint: disable=unnecessary-lambda
          for quantizer_id in sorted(non_weight_quantizers, key=lambda x: str(x)):
              activation_ctx = quantizer_id.ia_op_exec_context
              post_hooked_nx_node_key = nncf_graph.get_node_id_by_iap_context(activation_ctx)
              weight_quantizers = []
              for next_nx_node_key in nncf_graph.get_successors(post_hooked_nx_node_key):
                  weight_quantizers = traverse_graph(next_nx_node_key, weight_quantizers)
              if weight_quantizers:
                  activation_quantizer = self.non_weight_quantizers[quantizer_id].quantizer_module_ref
                  pairs.append((weight_quantizers, activation_quantizer))
          return pairs
  
      def _set_quantization_status(self, condition_fn: Callable[[BaseQuantizer], bool],
                                   apply_fn: Callable[[BaseQuantizer], None]):
          if self._model is not None:
              for m in self.all_quantizations.values():
                  if condition_fn(m):
                      apply_fn(m)
  
      def enable_activation_quantization(self):
          self._set_quantization_status(lambda x: not x.is_weights, lambda x: x.enable_quantization())
  
      def enable_weight_quantization(self):
          self._set_quantization_status(lambda x: x.is_weights, lambda x: x.enable_quantization())
  
      def disable_activation_quantization(self):
          self._set_quantization_status(lambda x: not x.is_weights, lambda x: x.disable_quantization())
  
      def disable_weight_quantization(self):
          self._set_quantization_status(lambda x: x.is_weights, lambda x: x.disable_quantization())
  
      def _get_local_init_range_config(self, scope: Scope, scope_overrides: Dict[str, Dict],
                                       global_init_range_config: Dict, quantizer_group: str):
          if isinstance(global_init_range_config, dict):
              module_init_range_config = global_init_range_config
          else:
              module_init_range_config = None
              matched_init_range_config = []
              for range_init_subconfig in global_init_range_config:
                  target_scopes = range_init_subconfig.get("target_scopes", None)
                  ignored_scopes = range_init_subconfig.get("ignored_scopes", None)
                  target_quantizer_group = range_init_subconfig.get("target_quantizer_group", quantizer_group)
                  if quantizer_group == target_quantizer_group and\
                       should_consider_scope(str(scope), target_scopes, ignored_scopes):
                      matched_init_range_config.append(range_init_subconfig)
  
              if len(matched_init_range_config) > 1:
                  raise AssertionError("The range initialization configs conflict with each other. "
                                       "Conflicting configs: {} for scope {}.".format(matched_init_range_config,
                                                                                      str(scope)))
  
  
              if len(matched_init_range_config) == 1:
                  module_init_range_config = matched_init_range_config[0]
              else:
                  raise AssertionError("The range initialization configs conflict with each other. "
                                       "Conflicting configs: {} for scope {}.".format(matched_init_range_config,
                                                                                      str(scope)))
  
          for overridden_scope in scope_overrides.keys():
              if in_scope_list(str(scope), overridden_scope):
                  override_config = scope_overrides[overridden_scope].get('initializer', {}).get("range")
                  if override_config is not None:
                      module_init_range_config = override_config
  
          if module_init_range_config is None:
              module_init_range_config = self.update_range_config_by_default({})
  
          return module_init_range_config
  
      def statistics(self):
          stats = super().statistics()
          num_enabled_quantization = len([1 for q in self.all_quantizations.values() if q.is_enabled_quantization()])
          multiplier = 100 / len(self.all_quantizations)
          stats["ratio_of_enabled_quantizations"] = num_enabled_quantization * multiplier
          if self._collect_compression_metrics:
              self.update_metric_store()
              for metric in self.metric_store.values():
                  for add_info, table in metric.items():
                      stats[add_info] = table
          return stats
  
  
  class QuantizationDebugInterface(DebugInterface):
      QUANTIZERS_IN_NNCF_MODULES_TRACKER_NAME = 'quantized_modules'
      ACTIVATION_QUANTIZERS_TRACKER_NAME = 'activation_quantizers'
      FUNCTION_QUANTIZERS_TRACKER_NAME = 'function_quantizers'
  
      def __init__(self):
          self.call_trackers = {
              self.QUANTIZERS_IN_NNCF_MODULES_TRACKER_NAME: CallCountTracker(
                  QuantizationDebugInterface.QUANTIZERS_IN_NNCF_MODULES_TRACKER_NAME),
              self.ACTIVATION_QUANTIZERS_TRACKER_NAME: CallCountTracker(
                  QuantizationDebugInterface.ACTIVATION_QUANTIZERS_TRACKER_NAME),
              self.FUNCTION_QUANTIZERS_TRACKER_NAME: CallCountTracker(
                  self.FUNCTION_QUANTIZERS_TRACKER_NAME)
          }
          self.graph_size = 0
  
          from nncf.debug import DEBUG_LOG_DIR
          self.dump_dir = Path(DEBUG_LOG_DIR) / Path("debug_dumps")
          self.dump_dir.mkdir(parents=True, exist_ok=True)
          self.scale_dump_dir = self.dump_dir / Path("scale")
          self.prop_graph_dump_dir = self.dump_dir / Path("quant_prop")
          if self.prop_graph_dump_dir.exists():
              shutil.rmtree(str(self.prop_graph_dump_dir))
          self.forward_call_count = 0
          self._strict_forward = False
  
      def init_actual(self, owner_model: NNCFNetwork):
          quantization_types = [class_type.__name__ for class_type in QUANTIZATION_MODULES.registry_dict.values()]
          quantizers_in_nncf_modules = owner_model.get_modules_in_nncf_modules_by_type(quantization_types)
          nncf_module_quantizations_id_list = [str(scope) for scope in
                                               quantizers_in_nncf_modules.keys()]  # type: List[str]
  
          activation_quantizer_id_list = owner_model.get_compression_modules_by_type(
              CompressionModuleType.ACTIVATION_QUANTIZER).keys()  # type: List[str]
          function_input_quantizer_id_list = owner_model.get_compression_modules_by_type(
              CompressionModuleType.FUNCTION_QUANTIZER).keys()  # type: List[str]
          self.call_trackers[self.QUANTIZERS_IN_NNCF_MODULES_TRACKER_NAME].init_with_key_list(
              nncf_module_quantizations_id_list)
          self.call_trackers[self.ACTIVATION_QUANTIZERS_TRACKER_NAME].init_with_key_list(
              activation_quantizer_id_list)
          self.call_trackers[self.FUNCTION_QUANTIZERS_TRACKER_NAME].init_with_key_list(
              function_input_quantizer_id_list)
          if self.scale_dump_dir.exists():
              shutil.rmtree(str(self.scale_dump_dir))
          self.scale_dump_dir.mkdir(parents=True, exist_ok=True)
          self._strict_forward = True
  
      def pre_forward_actions(self, module: 'NNCFNetwork'):
          self.reset_counters()
  
      def post_forward_actions(self, module: 'NNCFNetwork'):
          self.register_forward_call()
          # pylint:disable=protected-access
          ctx = module.get_tracing_context()
          self.set_graph_size(ctx.graph.get_nodes_count())
  
          quantization_types = [class_type.__name__ for class_type in QUANTIZATION_MODULES.registry_dict.values()]
          nncf_module_quantizations = module.get_modules_in_nncf_modules_by_type(
              quantization_types)  # type: Dict['Scope', nn.Module]
  
          for qm_scope, qm_module in nncf_module_quantizations.items():
              # Important - this will not work for DataParallel since it copies the
              # entire parent module for each thread and the `call_count` attributes
              # are incremented for thread local copies of `qm_module`, which are not
              # the same as the primary copies of `qm_module` iterated over at this point
              self.register_quantizer_module_call(str(qm_scope), qm_module.call_count)
              self.dump_scale(qm_module.get_trainable_params(), str(qm_scope))
              qm_module.reset_call_counter()
          self.print_call_stats()
  
          call_dict = ctx.get_node_call_counter_dict()
          total_calls = sum(call_dict.values())
          nncf_logger.debug("{} nodes called out of total {}".format(total_calls,
                                                                     ctx.graph.get_nodes_count()))
          if self._strict_forward:
              for tracker in self.call_trackers.values():
                  if tracker.get_never_called_keys():
                      # This will always trigger for DataParallel - disregard or disable debug mode
                      # for DataParallel runs
                      raise RuntimeError("{} has never called modules: {}!".format(
                          tracker.name, tracker.get_never_called_keys()))
  
      def dump_scale(self, quantizer_scale_params: Dict[str, torch.Tensor], quantizer_name: str):
          import re
          quantizer_normalized_name = re.sub(r'[^\w\-_\. ]', '_', quantizer_name)
          for scale_param_name, scale_param in quantizer_scale_params.items():
              fname = "{}_{}.txt".format(quantizer_normalized_name, scale_param_name)
              with open(str(self.scale_dump_dir / fname), "ba") as file:
                  np.savetxt(file, scale_param.cpu().numpy().flatten())
  
      def reset_counters(self):
          for tracker in self.call_trackers.values():
              tracker.reset()
  
      def register_quantizer_module_call(self, key, counts=None):
          self.call_trackers[self.QUANTIZERS_IN_NNCF_MODULES_TRACKER_NAME].register_call(key, counts)
  
      def register_activation_quantize_call(self, key: str):
          self.call_trackers[self.ACTIVATION_QUANTIZERS_TRACKER_NAME].register_call(key)
  
      def register_function_quantizer_call(self, key: str):
          self.call_trackers[self.FUNCTION_QUANTIZERS_TRACKER_NAME].register_call(key)
  
      def print_call_stats(self):
          nncf_logger.debug(" Graph size: {} nodes".format(self.graph_size))
          for tracker in self.call_trackers.values():
              msg = " {} tracker:".format(tracker.name)
              msg += " {} total calls;".format(tracker.get_total_call_count())
  
              never_called = tracker.get_never_called_keys()
              if never_called:
                  msg += " {} entries never called;".format(len(never_called))
  
              overcalled = tracker.get_overcalled_keys_with_call_counts()
              if overcalled:
                  msg += " {} entries called more than once;".format(len(overcalled))
              nncf_logger.debug(msg)
  
      def set_graph_size(self, new_size):
          if new_size != self.graph_size:
              nncf_logger.debug('\n')
              nncf_logger.debug(
                  " warning - graph size has changed from {} to {} since last forward".format(self.graph_size,
                                                                                              new_size))
          self.graph_size = new_size
  
      def register_forward_call(self):
          self.forward_call_count += 1
  
      def visualize_quantizer_propagation(self,
                                          prop_solver: QuantizerPropagationSolver,
                                          prop_graph: QuantizerPropagationStateGraph,
                                          iteration: str):
          self.prop_graph_dump_dir.mkdir(parents=True, exist_ok=True)
          fname = "quant_prop_iter_{}.dot".format(iteration)
          prop_solver.debug_visualize(prop_graph,
                                      self.prop_graph_dump_dir / Path(fname))
  
      def visualize_insertion_point_graph(self, insertion_point_graph: InsertionPointGraph):
          out_graph = nx.MultiDiGraph()
          for node_key, node in insertion_point_graph.nodes.items():
              if node[InsertionPointGraph.NODE_TYPE_NODE_ATTR] == InsertionPointGraphNodeType.INSERTION_POINT:
                  insertion_point_data = node[InsertionPointGraph.INSERTION_POINT_DATA_NODE_ATTR]  # type: InsertionPoint
                  label = "IP: {}".format(insertion_point_data.insertion_type)
                  out_graph.add_node(node_key, label=label, color="red")
              elif node[InsertionPointGraph.NODE_TYPE_NODE_ATTR] == InsertionPointGraphNodeType.OPERATOR:
                  out_graph.add_node(node_key)
              else:
                  raise RuntimeError("Invalid InsertionPointGraph node!")
          for u, v in insertion_point_graph.edges:
              out_graph.add_edge(u, v)
  
          for node_key, node in insertion_point_graph.nodes.items():
              if node[InsertionPointGraph.NODE_TYPE_NODE_ATTR] == InsertionPointGraphNodeType.OPERATOR:
                  for ip_node_key in node[InsertionPointGraph.ASSOCIATED_IP_NODE_KEYS_NODE_ATTR]:
                      out_graph.add_edge(node_key, ip_node_key, style="dashed", headport='e', tailport='e')
  
          nx.drawing.nx_pydot.write_dot(out_graph, self.dump_dir / Path("insertion_point_graph.dot"))
  