






from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import (
							    TF_MODEL_FOR_PRETRAINING_MAPPING,
							    TFMobileBertForMaskedLM,
							    TFMobileBertForMultipleChoice,
							    TFMobileBertForNextSentencePrediction,
							    TFMobileBertForPreTraining,
							    TFMobileBertForQuestionAnswering,
							    TFMobileBertForSequenceClassification,
							    TFMobileBertForTokenClassification,
							    TFMobileBertModel,
							)

@require_tf
class 				lowerCAmelCase_	(		__snake_case   ,		__snake_case   ,		unittest.TestCase			):
				_UpperCamelCase					:       Optional[int]							       = (
				    (
				        TFMobileBertModel,
				        TFMobileBertForMaskedLM,
				        TFMobileBertForNextSentencePrediction,
				        TFMobileBertForPreTraining,
				        TFMobileBertForQuestionAnswering,
				        TFMobileBertForSequenceClassification,
				        TFMobileBertForTokenClassification,
				        TFMobileBertForMultipleChoice,
				    )
				    if is_tf_available()
				    else ()
				)
				_UpperCamelCase					:       List[str]							       = (
				    {
				        "feature-extraction": TFMobileBertModel,
				        "fill-mask": TFMobileBertForMaskedLM,
				        "question-answering": TFMobileBertForQuestionAnswering,
				        "text-classification": TFMobileBertForSequenceClassification,
				        "token-classification": TFMobileBertForTokenClassification,
				        "zero-shot": TFMobileBertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				_UpperCamelCase					:       Any							       = False
				_UpperCamelCase					:       Dict							       = False


				def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase=False      ):
											_lowercase							:   Any 							=   super()._prepare_for_class(_lowerCAmelCase ,						_lowerCAmelCase ,						return_labels=_lowerCAmelCase      )

											if return_labels:
																		if model_class in get_values(_lowerCAmelCase      ):
																									_lowercase							:   str 							=   tf.zeros(self.model_tester.batch_size ,						dtype=tf.intaa      )

											return inputs_dict

				class 				lowerCAmelCase_	(		__snake_case			):


								def __init__(					self ,						_lowerCAmelCase ,						_lowerCAmelCase=1_3 ,						_lowerCAmelCase=7 ,						_lowerCAmelCase=True ,						_lowerCAmelCase=True ,						_lowerCAmelCase=True ,						_lowerCAmelCase=True ,						_lowerCAmelCase=9_9 ,						_lowerCAmelCase=3_2 ,						_lowerCAmelCase=3_2 ,						_lowerCAmelCase=2 ,						_lowerCAmelCase=4 ,						_lowerCAmelCase=3_7 ,						_lowerCAmelCase="gelu" ,						_lowerCAmelCase=0.1 ,						_lowerCAmelCase=0.1 ,						_lowerCAmelCase=5_1_2 ,						_lowerCAmelCase=1_6 ,						_lowerCAmelCase=2 ,						_lowerCAmelCase=0.02 ,						_lowerCAmelCase=3 ,						_lowerCAmelCase=4 ,						_lowerCAmelCase=None ,						):
															_lowercase							:   Optional[int] 							=   parent
															_lowercase							:   List[str] 							=   batch_size
															_lowercase							:   Any 							=   seq_length
															_lowercase							:   str 							=   is_training
															_lowercase							:   Optional[Any] 							=   use_input_mask
															_lowercase							:   Dict 							=   use_token_type_ids
															_lowercase							:   List[Any] 							=   use_labels
															_lowercase							:   Dict 							=   vocab_size
															_lowercase							:   Tuple 							=   hidden_size
															_lowercase							:   List[Any] 							=   num_hidden_layers
															_lowercase							:   List[Any] 							=   num_attention_heads
															_lowercase							:   Tuple 							=   intermediate_size
															_lowercase							:   Union[str, Any] 							=   hidden_act
															_lowercase							:   Dict 							=   hidden_dropout_prob
															_lowercase							:   List[Any] 							=   attention_probs_dropout_prob
															_lowercase							:   Optional[Any] 							=   max_position_embeddings
															_lowercase							:   List[Any] 							=   type_vocab_size
															_lowercase							:   Tuple 							=   type_sequence_label_size
															_lowercase							:   Optional[Any] 							=   initializer_range
															_lowercase							:   Tuple 							=   num_labels
															_lowercase							:   Any 							=   num_choices
															_lowercase							:   int 							=   scope
															_lowercase							:   Dict 							=   embedding_size


								def 						__a					(					self      ):
															_lowercase							:   str 							=   ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size      )

															_lowercase							:   Tuple 							=   None
															if self.use_input_mask:
																						_lowercase							:   List[str] 							=   random_attention_mask([self.batch_size, self.seq_length]      )

															_lowercase							:   Dict 							=   None
															if self.use_token_type_ids:
																						_lowercase							:   Any 							=   ids_tensor([self.batch_size, self.seq_length] ,						self.type_vocab_size      )

															_lowercase							:   Tuple 							=   None
															_lowercase							:   int 							=   None
															_lowercase							:   Optional[int] 							=   None
															if self.use_labels:
																						_lowercase							:   Any 							=   ids_tensor([self.batch_size] ,						self.type_sequence_label_size      )
																						_lowercase							:   str 							=   ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels      )
																						_lowercase							:   Optional[Any] 							=   ids_tensor([self.batch_size] ,						self.num_choices      )

															_lowercase							:   List[str] 							=   MobileBertConfig(
															    vocab_size=self.vocab_size ,						hidden_size=self.hidden_size ,						num_hidden_layers=self.num_hidden_layers ,						num_attention_heads=self.num_attention_heads ,						intermediate_size=self.intermediate_size ,						hidden_act=self.hidden_act ,						hidden_dropout_prob=self.hidden_dropout_prob ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob ,						max_position_embeddings=self.max_position_embeddings ,						type_vocab_size=self.type_vocab_size ,						initializer_range=self.initializer_range ,						embedding_size=self.embedding_size ,						)

															return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   Union[str, Any] 							=   TFMobileBertModel(config=_lowerCAmelCase      )
															_lowercase							:   Union[str, Any] 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   Optional[Any] 							=   model(_lowerCAmelCase      )

															_lowercase							:   Union[str, Any] 							=   [input_ids, input_mask]
															_lowercase							:   Optional[Any] 							=   model(_lowerCAmelCase      )

															_lowercase							:   List[Any] 							=   model(_lowerCAmelCase      )

															self.parent.assertEqual(
															    result.last_hidden_state.shape ,						(self.batch_size, self.seq_length, self.hidden_size)      )
															self.parent.assertEqual(result.pooler_output.shape ,						(self.batch_size, self.hidden_size)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   int 							=   TFMobileBertForMaskedLM(config=_lowerCAmelCase      )
															_lowercase							:   List[str] 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   Optional[int] 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   List[str] 							=   TFMobileBertForNextSentencePrediction(config=_lowerCAmelCase      )
															_lowercase							:   Dict 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   str 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(result.logits.shape ,						(self.batch_size, 2)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   Tuple 							=   TFMobileBertForPreTraining(config=_lowerCAmelCase      )
															_lowercase							:   Dict 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   List[str] 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(
															    result.prediction_logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size)      )
															self.parent.assertEqual(result.seq_relationship_logits.shape ,						(self.batch_size, 2)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   str 							=   self.num_labels
															_lowercase							:   Union[str, Any] 							=   TFMobileBertForSequenceClassification(config=_lowerCAmelCase      )
															_lowercase							:   Any 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   Optional[int] 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_labels)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   Any 							=   self.num_choices
															_lowercase							:   Optional[int] 							=   TFMobileBertForMultipleChoice(config=_lowerCAmelCase      )
															_lowercase							:   str 							=   tf.tile(tf.expand_dims(_lowerCAmelCase ,						1      ) ,						(1, self.num_choices, 1)      )
															_lowercase							:   int 							=   tf.tile(tf.expand_dims(_lowerCAmelCase ,						1      ) ,						(1, self.num_choices, 1)      )
															_lowercase							:   Tuple 							=   tf.tile(tf.expand_dims(_lowerCAmelCase ,						1      ) ,						(1, self.num_choices, 1)      )
															_lowercase							:   Dict 							=   {
															    'input_ids': multiple_choice_inputs_ids,
															    'attention_mask': multiple_choice_input_mask,
															    'token_type_ids': multiple_choice_token_type_ids,
															}
															_lowercase							:   List[Any] 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_choices)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   Union[str, Any] 							=   self.num_labels
															_lowercase							:   Tuple 							=   TFMobileBertForTokenClassification(config=_lowerCAmelCase      )
															_lowercase							:   Optional[int] 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   Union[str, Any] 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels)      )


								def 						__a					(					self ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase ,						_lowerCAmelCase      ):
															_lowercase							:   Tuple 							=   TFMobileBertForQuestionAnswering(config=_lowerCAmelCase      )
															_lowercase							:   Optional[int] 							=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
															_lowercase							:   List[str] 							=   model(_lowerCAmelCase      )
															self.parent.assertEqual(result.start_logits.shape ,						(self.batch_size, self.seq_length)      )
															self.parent.assertEqual(result.end_logits.shape ,						(self.batch_size, self.seq_length)      )


								def 						__a					(					self      ):
															_lowercase							:   str 							=   self.prepare_config_and_inputs()
															(
															    (
															    _lowercase
															)							,							(
															    _lowercase
															)							,							(
															    _lowercase
															)							,							(
															    _lowercase
															)							,							(
															    _lowercase
															)							,							(
															    _lowercase
															)							,							(
															    _lowercase
															)							,							
															)							:   Tuple 							=   config_and_inputs
															_lowercase							:   Dict 							=   {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
															return config, inputs_dict




				def 						__a					(					self      ):
											_lowercase							:   str 							=   TFMobileBertModelTest.TFMobileBertModelTester(self      )
											_lowercase							:   Optional[int] 							=   ConfigTester(self ,						config_class=_lowerCAmelCase ,						hidden_size=3_7      )


				def 						__a					(					self      ):
											self.config_tester.run_common_tests()


				def 						__a					(					self      ):
											_lowercase							:   int 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_model(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   Tuple 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_masked_lm(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   Dict 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_multiple_choice(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   str 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   List[str] 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_pretraining(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   str 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_question_answering(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   Dict 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_sequence_classification(*_lowerCAmelCase      )


				def 						__a					(					self      ):
											_lowercase							:   Optional[Any] 							=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_mobilebert_for_token_classification(*_lowerCAmelCase      )


				@slow
				def 						__a					(					self      ):
											# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											for model_name in ["google/mobilebert-uncased"]:
																		_lowercase							:   List[str] 							=   TFMobileBertModel.from_pretrained(_lowerCAmelCase      )
																		self.assertIsNotNone(_lowerCAmelCase      )




@require_tf
class 				lowerCAmelCase_	(		unittest.TestCase			):


				@slow
				def 						__a					(					self      ):
											_lowercase							:   Any 							=   TFMobileBertForPreTraining.from_pretrained('google/mobilebert-uncased'      )
											_lowercase							:   int 							=   tf.constant([[0, 1, 2, 3, 4, 5]]      )
											_lowercase							:   Tuple 							=   model(_lowerCAmelCase      )[0]

											_lowercase							:   str 							=   [1, 6, 3_0_5_2_2]
											self.assertEqual(output.shape ,						_lowerCAmelCase      )

											_lowercase							:   int 							=   tf.constant(
											    [
											        [
											            [-4.5_91_95_47, -9.24_82_95, -9.64_52_56],
											            [-6.7_30_61_75, -6.44_02_84, -6.6_05_28_37],
											            [-7.2_74_35_06, -6.7_84_79_15, -6.02_46_73],
											        ]
											    ]      )
											tf.debugging.assert_near(output[:, :3, :3] ,						_lowerCAmelCase ,						atol=1E-4      )





'''simple docstring'''

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


UpperCamelCase__  :						Optional[Any]					=      logging.getLogger(__name__)

UpperCamelCase__  :						Dict					=      50  # max width of layer names
UpperCamelCase__  :						Any					=      70  # max width of quantizer names







def __UpperCamelCase( _A       :					Union[str, Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					List[str] 					=	parser.add_argument_group('''quant_trainer arguments'''     )
   group.add_argument('''--wprec'''	,      type=_A	,      default=8	,      help='''weight precision'''     )
   group.add_argument('''--aprec'''	,      type=_A	,      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=_A	,      nargs='''+'''	,      help='''disable quantizers by keyword'''     )
   group.add_argument('''--quant-disable-layer-module'''	,      type=_A	,      help='''disable quantizers by keyword under layer.'''     )
   group.add_argument('''--quant-enable-layer-module'''	,      type=_A	,      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=_A	,      type=_A	,      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=_A	,      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( _A       :					Tuple     ):
   '''simple docstring'''

   if args.calibrator == "max":
      UpperCAmelCase__				:					str 					=	'''max'''
   elif args.calibrator == "percentile":
      if args.percentile is None:
         raise ValueError('''Specify --percentile when using percentile calibrator'''     )
      UpperCAmelCase__				:					Dict 					=	'''histogram'''
   elif args.calibrator == "mse":
      UpperCAmelCase__				:					Any 					=	'''histogram'''
   else:
      raise ValueError(F'''Invalid calibrator {args.calibrator}'''     )

   UpperCAmelCase__				:					Dict 					=	QuantDescriptor(num_bits=args.aprec	,      calib_method=_A     )
   UpperCAmelCase__				:					str 					=	QuantDescriptor(num_bits=args.wprec	,      axis=(None if args.quant_per_tensor else (0,))     )
   quant_nn.QuantLinear.set_default_quant_desc_input(_A     )
   quant_nn.QuantLinear.set_default_quant_desc_weight(_A     )







def __UpperCamelCase( _A       :					Any	,      _A       :					Any	,      _A       :					Any=False	,      _A       :					Optional[Any]=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(_A	,      ['''embeddings''']	,      which='''weight'''	,      _disabled=_A     )

      if args.quant_disable:
         set_quantizer_by_name(_A	,      ['''''']	,      _disabled=_A     )

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

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

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

      if args.recalibrate_weights:
         recalibrate_weights(_A     )

      if args.fuse_qkv:
         fuse_qkv(_A	,      _A     )

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

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







def __UpperCamelCase( _A       :					str     ):
   '''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( _A       :					Tuple	,      _A       :					Any     ):
   '''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(_A     )







def __UpperCamelCase( _A       :					Dict	,      _A       :					Optional[int]     ):
   '''simple docstring'''

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

      UpperCAmelCase__				:					Dict 					=	max(_A	,      _A	,      _A     )
      qq._amax.fill_(_A     )
      qk._amax.fill_(_A     )
      qv._amax.fill_(_A     )
      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( _A       :					Dict	,      _A       :					Any     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if name.endswith('''.output.dense'''     ) and not name.endswith('''attention.output.dense'''     ):
         UpperCAmelCase__				:					Union[str, Any] 					=	mod._input_quantizer._amax.data.detach().item()
         mod._input_quantizer._amax.data.detach().clamp_(max=_A     )
         UpperCAmelCase__				:					Tuple 					=	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( _A       :					str     ):
   '''simple docstring'''

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







def __UpperCamelCase( _A       :					List[str]     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if hasattr(_A	,      '''_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)
         UpperCAmelCase__				:					Any 					=	set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis     )
         UpperCAmelCase__				:					Optional[Any] 					=	set(range(len(mod.weight.size()     )     )     ) - axis_set
         UpperCAmelCase__				:					int 					=	pytorch_quantization.utils.reduce_amax(mod.weight	,      axis=_A	,      keepdims=_A     ).detach()
         logger.info(F'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}'''     )
         UpperCAmelCase__				:					str 					=	amax







def __UpperCamelCase( _A       :					Dict	,      _A       :					Tuple=25	,      _A       :					Any=1_80	,      _A       :					Optional[int]=None     ):
   '''simple docstring'''

   if ignore is None:
      UpperCAmelCase__				:					Dict 					=	[]
   elif not isinstance(_A	,      _A     ):
      UpperCAmelCase__				:					int 					=	[ignore]

   UpperCAmelCase__				:					Optional[int] 					=	0
   for name, mod in model.named_modules():
      if not hasattr(_A	,      '''weight'''     ):
         continue
      UpperCAmelCase__				:					Dict 					=	max(_A	,      len(_A     )     )

   for name, mod in model.named_modules():
      UpperCAmelCase__				:					str 					=	getattr(_A	,      '''_input_quantizer'''	,      _A     )
      UpperCAmelCase__				:					int 					=	getattr(_A	,      '''_weight_quantizer'''	,      _A     )
      if not hasattr(_A	,      '''weight'''     ):
         continue
      if type(_A     ) in ignore:
         continue
      if [True for s in ignore if type(_A     ) is str and s in name]:
         continue
      UpperCAmelCase__				:					Dict 					=	F'''Act:{input_q.extra_repr()}'''
      UpperCAmelCase__				:					int 					=	F'''Wgt:{weight_q.extra_repr()}'''
      UpperCAmelCase__				:					Dict 					=	F'''{name:{name_width}} {act_str} {wgt_str}'''
      if len(_A     ) <= line_width:
         logger.info(_A     )
      else:
         logger.info(F'''{name:{name_width}} {act_str}'''     )
         logger.info(F'''{'  ':{name_width}} {wgt_str}'''     )







def __UpperCamelCase( _A       :					Optional[Any]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					str 					=	0
   for name, mod in model.named_modules():
      if isinstance(_A	,      pytorch_quantization.nn.TensorQuantizer     ):
         print(F'''{name:80} {mod}'''     )
         count += 1
   print(F'''{count} TensorQuantizers found in model'''     )







def __UpperCamelCase( _A       :					Dict	,      _A       :					Optional[Any]	,      _A       :					Union[str, Any]	,      _A       :					Union[str, Any]	,      _A       :					Dict     ):
   '''simple docstring'''
   UpperCAmelCase__				:					Any 					=	getattr(_A	,      _A	,      _A     )
   if quantizer_mod is not None:
      assert hasattr(_A	,      _A     )
      setattr(_A	,      _A	,      _A     )
   else:
      logger.warning(F'''{name} has no {quantizer}'''     )







def __UpperCamelCase( _A       :					str	,      _A       :					Any	,      _A       :					Optional[int]="both"	,      **_A       :					List[str]     ):
   '''simple docstring'''
   UpperCAmelCase__				:					Tuple 					=	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(_A	,      _A	,      '''_input_quantizer'''	,      _A	,      _A     )
      if which in ["weight", "both"]:
         set_quantizer(_A	,      _A	,      '''_weight_quantizer'''	,      _A	,      _A     )
   logger.info(_A     )







def __UpperCamelCase( _A       :					Tuple	,      _A       :					List[str]	,      **_A       :					Optional[Any]     ):
   '''simple docstring'''

   for name, mod in model.named_modules():
      if hasattr(_A	,      '''_input_quantizer'''     ) or hasattr(_A	,      '''_weight_quantizer'''     ):
         for n in names:
            if re.search(_A	,      _A     ):
               set_quantizers(_A	,      _A	,      **_A     )
      elif name.endswith('''_quantizer'''     ):
         for n in names:
            if re.search(_A	,      _A     ):
               UpperCAmelCase__				:					str 					=	F'''Warning: changing {name:{name_width}}'''
               for k, v in kwargs.items():
                  s += F''' {k}={v}'''
                  setattr(_A	,      _A	,      _A     )
               logger.info(_A     )




