docstr: Two changes will be done during initialization. One is an attribute named ``_nni_wrapper`` will be set to original module,
  this attribute points to this wrapper in the original module.
  The other is the original ``module.forward`` will be replaced by ``module._nni_wrapper.forward``.
  
  The module can be unwrapped by ``module._nni_wrapper.unwrap()``.
  
  Parameters
  ----------
  module
      The torch.nn.Module to be wrapped.
  module_name
      The name of the module in the original model.
  config
      The config is a dict which contains keys (not required): ``pruning``, ``quantization``, ``distillation``.
  fused_modules:
      The List contains a series of modules which need to fuse.
docstr_tok: ["Two","changes","will","be","done","during","initialization",".","One","is","an","attribute","named","``","_nni_wrapper","``","will","be","set","to","original","module",",","this","attribute","points","to","this","wrapper","in","the","original","module",".","The","other","is","the","original","``","module.forward","``","will","be","replaced","by","``","module._nni_wrapper.forward","``",".","The","module","can","be","unwrapped","by","``","module._nni_wrapper.unwrap","(",")","``",".","Parameters","--","--","--","--","--","module","The","torch.nn.Module","to","be","wrapped",".","module_name","The","name","of","the","module","in","the","original","model",".","config","The","config","is","a","dict","which","contains","keys","(","not","required",")",":","``","pruning","``",",","``","quantization","``",",","``","distillation","``",".","fused_modules",":","The","List","contains","a","series","of","modules","which","need","to","fuse","."]
code: def __init__(
      self,
      module: torch.nn.Module,
      module_name: str,
      config: Dict[str, Dict[str, Any]] | None = None,
      fused_modules: List[torch.nn.Module] | None = None,
  ):
      """
      Two changes will be done during initialization. One is an attribute named ``_nni_wrapper`` will be set to original module,
      this attribute points to this wrapper in the original module.
      The other is the original ``module.forward`` will be replaced by ``module._nni_wrapper.forward``.
  
      The module can be unwrapped by ``module._nni_wrapper.unwrap()``.
  
      Parameters
      ----------
      module
          The torch.nn.Module to be wrapped.
      module_name
          The name of the module in the original model.
      config
          The config is a dict which contains keys (not required): ``pruning``, ``quantization``, ``distillation``.
      fused_modules:
          The List contains a series of modules which need to fuse.
      """
      super().__init__()
  
      # origin layer information
      assert isinstance(module, torch.nn.Module)
      object.__setattr__(self, "module", module)
      self.module: torch.nn.Module
      self.module_forward = self.module.forward
      self.name = module_name
      self.config = config if config is not None else {}
      # config storage
      self.configs = {}
      for mode, value in self.config.items():
          if mode not in self.configs:
              self.configs[mode] = []
          self.configs[mode].append(value)
  
      assert all(
          k in ["pruning", "quantization", "distillation"]
          for k in self.config
      )
  
      # the arguments' name of self.module.forward
      self._input_args_spec = inspect.getfullargspec(
          self.module.forward
      )
  
      # create target spaces
      self.pruning_target_spaces: Dict[str, PruningTargetSpace] = {}
      self.quantization_target_spaces: Dict[
          str, QuantizationTargetSpace
      ] = {}
      self.distillation_target_spaces: Dict[
          str, DistillationTargetSpace
      ] = {}
  
      if "pruning" in self.config:
          self.extend_target_spaces(self.config["pruning"], "pruning")
      if "quantization" in self.config:
          self.extend_target_spaces(
              self.config["quantization"], "quantization"
          )
      if "distillation" in self.config:
          self.extend_target_spaces(
              self.config["distillation"], "distillation"
          )
  
      self._frozen = False
      # By default, input/output shape will be track during forward,
      # more track functions can be registered by ``ModuleWrapper.register_track_info_func``.
      # An example please refer ``track_target_shape``.
      self._track_funcs: List[
          Callable[[ModuleWrapper, str, Tensor], None]
      ] = [track_target_shape]
      self.fused_modules = (
          fused_modules if fused_modules is not None else []
      )
  
imports: from __future__ import annotations
  from collections import defaultdict
  import logging
  import inspect
  import functools
  from typing import Any, Callable, Dict, List, Tuple, Type, Union, Literal
  import torch
  import torch.nn
  from torch import Tensor
  from .apply_method import pruning_apply_methods, quant_apply_methods
  from .config import select_modules_by_config
  from .setting import INPUT_PREFIX, OUTPUT_PREFIX, canonicalize_settings
  from .target_space import TargetSpace, TargetType, PruningTargetSpace, QuantizationTargetSpace, DistillationTargetSpace
  from .fuse_modules import fuse_modules
  from ..utils import check_bias, update_config, validate_fused_modules_config, get_fused_module_list, get_nested_attr
label_desc: Use image node_id 1 to create a new ModuleWrapper object from inherited base classes: torch.nn with example: obj = ModuleWrapper(module, module_name, config, fused_modules)