






import logging
import re

import pytorch_quantization
import pytorch_quantization.nn as quant_nn
import torch
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor


lowerCamelCase__       		=  logging.getLogger(__name__)

lowerCamelCase__       		=  50  # max width of layer names
lowerCamelCase__       		=  70  # max width of quantizer names






def 	UpperCamelCase	(							snake_case__  : List[Any]	):




     '''simple docstring'''






     __snake_case      :str 							=    parser.add_argument_group("""quant_trainer arguments"""	)
     group.add_argument("""--wprec"""						,type=snake_case__						,default=8						,help="""weight precision"""	)
     group.add_argument("""--aprec"""						,type=snake_case__						,default=8						,help="""activation precision"""	)
     group.add_argument("""--quant-per-tensor"""						,action="""store_true"""						,help="""per tensor weight scaling"""	)
     group.add_argument("""--quant-disable"""						,action="""store_true"""						,help="""disable all quantizers"""	)
     group.add_argument("""--quant-disable-embeddings"""						,action="""store_true"""						,help="""disable all embeddings quantizers"""	)
     group.add_argument("""--quant-disable-keyword"""						,type=snake_case__						,nargs="""+"""						,help="""disable quantizers by keyword"""	)
     group.add_argument("""--quant-disable-layer-module"""						,type=snake_case__						,help="""disable quantizers by keyword under layer."""	)
     group.add_argument("""--quant-enable-layer-module"""						,type=snake_case__						,help="""enable quantizers by keyword under layer"""	)
     group.add_argument("""--calibrator"""						,default="""max"""						,help="""which quantization range calibrator to use"""	)
     group.add_argument("""--percentile"""						,default=snake_case__						,type=snake_case__						,help="""percentile for PercentileCalibrator"""	)
     group.add_argument("""--fuse-qkv"""						,action="""store_true"""						,help="""use the same scale factor for qkv"""	)
     group.add_argument("""--clip-gelu"""						,metavar="""N"""						,type=snake_case__						,help="""clip gelu output maximum value to N"""	)
     group.add_argument(
         """--recalibrate-weights"""						,action="""store_true"""						,help=(
             """recalibrate weight amaxes by taking the max of the weights."""
             """ amaxes will be computed with the current quantization granularity (axis)."""
         )						,)






def 	UpperCamelCase	(							snake_case__  : Optional[int]	):




     '''simple docstring'''






     if args.calibrator == "max":
          __snake_case      :List[Any] 							=    """max"""
     elif args.calibrator == "percentile":
          if args.percentile is None:
               raise ValueError("""Specify --percentile when using percentile calibrator"""	)
          __snake_case      :Any 							=    """histogram"""
     elif args.calibrator == "mse":
          __snake_case      :Tuple 							=    """histogram"""
     else:
          raise ValueError(f'''Invalid calibrator {args.calibrator}'''	)

     __snake_case      :Optional[Any] 							=    QuantDescriptor(num_bits=args.aprec						,calib_method=snake_case__	)
     __snake_case      :Tuple 							=    QuantDescriptor(num_bits=args.wprec						,axis=(None if args.quant_per_tensor else (0,))	)
     quant_nn.QuantLinear.set_default_quant_desc_input(snake_case__	)
     quant_nn.QuantLinear.set_default_quant_desc_weight(snake_case__	)






def 	UpperCamelCase	(							snake_case__  : List[str]						,snake_case__  : int						,snake_case__  : List[str]=False						,snake_case__  : Optional[int]=False	):




     '''simple docstring'''






     logger.info("""Configuring Model for Quantization"""	)
     logger.info(f'''using quantization package {pytorch_quantization.__file__}'''	)

     if not calib:
          if args.quant_disable_embeddings:
               set_quantizer_by_name(snake_case__						,["""embeddings"""]						,which="""weight"""						,_disabled=snake_case__	)

          if args.quant_disable:
               set_quantizer_by_name(snake_case__						,[""""""]						,_disabled=snake_case__	)

          if args.quant_disable_keyword:
               set_quantizer_by_name(snake_case__						,args.quant_disable_keyword						,_disabled=snake_case__	)

          if args.quant_disable_layer_module:
               set_quantizer_by_name(snake_case__						,[R"""layer.\d+.""" + args.quant_disable_layer_module]						,_disabled=snake_case__	)

          if args.quant_enable_layer_module:
               set_quantizer_by_name(snake_case__						,[R"""layer.\d+.""" + args.quant_enable_layer_module]						,_disabled=snake_case__	)

          if args.recalibrate_weights:
               recalibrate_weights(snake_case__	)

          if args.fuse_qkv:
               fuse_qkv(snake_case__						,snake_case__	)

     if args.clip_gelu:
          clip_gelu(snake_case__						,args.clip_gelu	)

     # if args.local_rank in [-1, 0] and not calib:
     print_quant_summary(snake_case__	)






def 	UpperCamelCase	(							snake_case__  : Union[str, Any]	):




     '''simple docstring'''






     logger.info("""Enabling Calibration"""	)
     for name, module in model.named_modules():
          if name.endswith("""_quantizer"""	):
               if module._calibrator is not None:
                    module.disable_quant()
                    module.enable_calib()
               else:
                    module.disable()
               logger.info(f'''{name:80}: {module}'''	)






def 	UpperCamelCase	(							snake_case__  : Union[str, Any]						,snake_case__  : List[str]	):




     '''simple docstring'''






     logger.info("""Loading calibrated amax"""	)
     for name, module in model.named_modules():
          if name.endswith("""_quantizer"""	):
               if module._calibrator is not None:
                    if isinstance(module._calibrator						,calib.MaxCalibrator	):
                         module.load_calib_amax()
                    else:
                         module.load_calib_amax("""percentile"""						,percentile=args.percentile	)
                    module.enable_quant()
                    module.disable_calib()
               else:
                    module.enable()
     model.cuda()
     print_quant_summary(snake_case__	)






def 	UpperCamelCase	(							snake_case__  : Any						,snake_case__  : Any	):




     '''simple docstring'''






     def fusea(snake_case__  : Optional[int]						,snake_case__  : Union[str, Any]						,snake_case__  : str	):
          for mod in [qq, qk, qv]:
               if not hasattr(snake_case__						,"""_amax"""	):
                    print("""          WARNING: NO AMAX BUFFER"""	)
                    return
          __snake_case      :List[Any] 							=    qq._amax.detach().item()
          __snake_case      :List[str] 							=    qk._amax.detach().item()
          __snake_case      :Dict 							=    qv._amax.detach().item()

          __snake_case      :int 							=    max(snake_case__						,snake_case__						,snake_case__	)
          qq._amax.fill_(snake_case__	)
          qk._amax.fill_(snake_case__	)
          qv._amax.fill_(snake_case__	)
          logger.info(f'''          q={q:5.2f} k={k:5.2f} v={v:5.2f} -> {amax:5.2f}'''	)

     for name, mod in model.named_modules():
          if name.endswith(""".attention.self"""	):
               logger.info(f'''FUSE_QKV: {name:{name_width}}'''	)
               fusea(mod.matmul_q_input_quantizer						,mod.matmul_k_input_quantizer						,mod.matmul_v_input_quantizer	)
               if args.quant_per_tensor:
                    fusea(mod.query._weight_quantizer						,mod.key._weight_quantizer						,mod.value._weight_quantizer	)






def 	UpperCamelCase	(							snake_case__  : Optional[int]						,snake_case__  : Any	):




     '''simple docstring'''






     for name, mod in model.named_modules():
          if name.endswith(""".output.dense"""	) and not name.endswith("""attention.output.dense"""	):
               __snake_case      :Union[str, Any] 							=    mod._input_quantizer._amax.data.detach().item()
               mod._input_quantizer._amax.data.detach().clamp_(max=snake_case__	)
               __snake_case      :List[str] 							=    mod._input_quantizer._amax.data.detach().item()
               logger.info(f'''CLIP_GELU: {name:{name_width}} amax: {amax_init:5.2f} -> {amax:5.2f}'''	)






def 	UpperCamelCase	(							snake_case__  : Optional[Any]	):




     '''simple docstring'''






     for name, mod in model.named_modules():
          if hasattr(snake_case__						,"""_weight_quantizer"""	) and mod._weight_quantizer.axis is not None:
               __snake_case      :Optional[Any] 							=    mod.weight.shape[0]
               __snake_case      :Optional[int] 							=    mod._weight_quantizer._amax.detach()
               __snake_case      :int 							=    torch.ones(snake_case__						,dtype=amax.dtype						,device=amax.device	) * amax
               print(f'''expanding {name} {amax} -> {mod._weight_quantizer._amax}'''	)






def 	UpperCamelCase	(							snake_case__  : Tuple	):




     '''simple docstring'''






     for name, mod in model.named_modules():
          if hasattr(snake_case__						,"""_weight_quantizer"""	):
               if not hasattr(mod.weight_quantizer						,"""_amax"""	):
                    print("""RECALIB: {name:{name_width}} WARNING: NO AMAX BUFFER"""	)
                    continue

               # determine which axes to reduce across
               # e.g. a 4D tensor quantized per axis 0 should reduce over (1,2,3)
               __snake_case      :Optional[int] 							=    set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis	)
               __snake_case      :Optional[Any] 							=    set(range(len(mod.weight.size()	)	)	) - axis_set
               __snake_case      :int 							=    pytorch_quantization.utils.reduce_amax(mod.weight						,axis=snake_case__						,keepdims=snake_case__	).detach()
               logger.info(f'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}'''	)
               __snake_case      :Union[str, Any] 							=    amax






def 	UpperCamelCase	(							snake_case__  : Union[str, Any]						,snake_case__  : Optional[int]=25						,snake_case__  : List[Any]=180						,snake_case__  : int=None	):




     '''simple docstring'''






     if ignore is None:
          __snake_case      :Optional[Any] 							=    []
     elif not isinstance(snake_case__						,snake_case__	):
          __snake_case      :List[Any] 							=    [ignore]

     __snake_case      :Any 							=    0
     for name, mod in model.named_modules():
          if not hasattr(snake_case__						,"""weight"""	):
               continue
          __snake_case      :int 							=    max(snake_case__						,len(snake_case__	)	)

     for name, mod in model.named_modules():
          __snake_case      :List[Any] 							=    getattr(snake_case__						,"""_input_quantizer"""						,snake_case__	)
          __snake_case      :Union[str, Any] 							=    getattr(snake_case__						,"""_weight_quantizer"""						,snake_case__	)
          if not hasattr(snake_case__						,"""weight"""	):
               continue
          if type(snake_case__	) in ignore:
               continue
          if [True for s in ignore if type(snake_case__	) is str and s in name]:
               continue
          __snake_case      :Dict 							=    f'''Act:{input_q.extra_repr()}'''
          __snake_case      :List[str] 							=    f'''Wgt:{weight_q.extra_repr()}'''
          __snake_case      :Tuple 							=    f'''{name:{name_width}} {act_str} {wgt_str}'''
          if len(snake_case__	) <= line_width:
               logger.info(snake_case__	)
          else:
               logger.info(f'''{name:{name_width}} {act_str}'''	)
               logger.info(f'''{'  ':{name_width}} {wgt_str}'''	)






def 	UpperCamelCase	(							snake_case__  : Tuple	):




     '''simple docstring'''






     __snake_case      :Optional[Any] 							=    0
     for name, mod in model.named_modules():
          if isinstance(snake_case__						,pytorch_quantization.nn.TensorQuantizer	):
               print(f'''{name:80} {mod}'''	)
               count += 1
     print(f'''{count} TensorQuantizers found in model'''	)






def 	UpperCamelCase	(							snake_case__  : Any						,snake_case__  : Optional[Any]						,snake_case__  : List[str]						,snake_case__  : Tuple						,snake_case__  : Optional[Any]	):




     '''simple docstring'''






     __snake_case      :List[Any] 							=    getattr(snake_case__						,snake_case__						,snake_case__	)
     if quantizer_mod is not None:
          assert hasattr(snake_case__						,snake_case__	)
          setattr(snake_case__						,snake_case__						,snake_case__	)
     else:
          logger.warning(f'''{name} has no {quantizer}'''	)






def 	UpperCamelCase	(							snake_case__  : str						,snake_case__  : List[Any]						,snake_case__  : Dict="both"						,**snake_case__  : Union[str, Any]	):




     '''simple docstring'''






     __snake_case      :List[str] 							=    f'''Warning: changing {which} quantizers of {name:{qname_width}}'''
     for k, v in kwargs.items():
          s += f''' {k}={v}'''
          if which in ["input", "both"]:
               set_quantizer(snake_case__						,snake_case__						,"""_input_quantizer"""						,snake_case__						,snake_case__	)
          if which in ["weight", "both"]:
               set_quantizer(snake_case__						,snake_case__						,"""_weight_quantizer"""						,snake_case__						,snake_case__	)
     logger.info(snake_case__	)






def 	UpperCamelCase	(							snake_case__  : Tuple						,snake_case__  : List[str]						,**snake_case__  : Tuple	):




     '''simple docstring'''






     for name, mod in model.named_modules():
          if hasattr(snake_case__						,"""_input_quantizer"""	) or hasattr(snake_case__						,"""_weight_quantizer"""	):
               for n in names:
                    if re.search(snake_case__						,snake_case__	):
                         set_quantizers(snake_case__						,snake_case__						,**snake_case__	)
          elif name.endswith("""_quantizer"""	):
               for n in names:
                    if re.search(snake_case__						,snake_case__	):
                         __snake_case      :List[Any] 							=    f'''Warning: changing {name:{name_width}}'''
                         for k, v in kwargs.items():
                              s += f''' {k}={v}'''
                              setattr(snake_case__						,snake_case__						,snake_case__	)
                         logger.info(snake_case__	)










from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


lowerCamelCase__       		=  logging.get_logger(__name__)

lowerCamelCase__       		=  {
    """microsoft/focalnet-tiny""": """https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json""",
}




class       snake_case__ (				lowercase_						,     lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					List[Any]								=			"focalnet"



       def __init__(				self      ,    a__=2_24      ,    a__=4      ,    a__=3      ,    a__=96      ,    a__=False      ,    a__=[1_92, 3_84, 7_68, 7_68]      ,    a__=[2, 2, 6, 2]      ,    a__=[2, 2, 2, 2]      ,    a__=[3, 3, 3, 3]      ,    a__="gelu"      ,    a__=4.0      ,    a__=0.0      ,    a__=0.1      ,    a__=False      ,    a__=1e-4      ,    a__=False      ,    a__=False      ,    a__=False      ,    a__=0.02      ,    a__=1e-5      ,    a__=32      ,    a__=None      ,    a__=None      ,    **a__      ,    )			->   Tuple:






            '''simple docstring'''





            super().__init__(**a__     )

            __snake_case      :Optional[Any] 							=    image_size
            __snake_case      :Dict 							=    patch_size
            __snake_case      :Tuple 							=    num_channels
            __snake_case      :List[str] 							=    embed_dim
            __snake_case      :Any 							=    use_conv_embed
            __snake_case      :List[Any] 							=    hidden_sizes
            __snake_case      :str 							=    depths
            __snake_case      :List[Any] 							=    focal_levels
            __snake_case      :int 							=    focal_windows
            __snake_case      :Any 							=    hidden_act
            __snake_case      :List[str] 							=    mlp_ratio
            __snake_case      :str 							=    hidden_dropout_prob
            __snake_case      :List[str] 							=    drop_path_rate
            __snake_case      :int 							=    use_layerscale
            __snake_case      :int 							=    layerscale_value
            __snake_case      :Tuple 							=    use_post_layernorm
            __snake_case      :Dict 							=    use_post_layernorm_in_modulation
            __snake_case      :Union[str, Any] 							=    normalize_modulator
            __snake_case      :List[str] 							=    initializer_range
            __snake_case      :Optional[Any] 							=    layer_norm_eps
            __snake_case      :Tuple 							=    encoder_stride
            __snake_case      :List[str] 							=    ["""stem"""] + [F'''stage{idx}''' for idx in range(1      ,    len(self.depths     ) + 1     )]
            __snake_case							,			__snake_case      :str 							=    get_aligned_output_features_output_indices(
                out_features=a__      ,    out_indices=a__      ,    stage_names=self.stage_names     )


