





import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.generation import (
         FlaxForcedBOSTokenLogitsProcessor,
         FlaxForcedEOSTokenLogitsProcessor,
         FlaxLogitsProcessorList,
         FlaxMinLengthLogitsProcessor,
         FlaxTemperatureLogitsWarper,
         FlaxTopKLogitsWarper,
         FlaxTopPLogitsWarper,
     )

@require_flax
class   A							(						unittest.TestCase    ):



      def 						lowerCAmelCase__     (						self:    Dict		,	_lowerCAmelCase:    int		,	_lowerCAmelCase:    int							)    ->     List[str]:
         '''simple docstring'''

         UpperCAmelCase_           =jnp.ones((batch_size, length)							) / length
         return scores



      def 						lowerCAmelCase__     (						self:    List[str]							)    ->     Optional[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =None
         UpperCAmelCase_           =20

         UpperCAmelCase_           =self._get_uniform_logits(batch_size=2		,	length=_lowerCAmelCase							)

         # tweak scores to not be uniform anymore
         UpperCAmelCase_           =scores.at[1, 5].set((1 / length) + 0.1							)  # peak, 1st batch
         UpperCAmelCase_           =scores.at[1, 10].set((1 / length) - 0.4							)  # valley, 1st batch

         # compute softmax
         UpperCAmelCase_           =jax.nn.softmax(_lowerCAmelCase		,	axis=-1							)

         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=0.5							)
         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=1.3							)

         UpperCAmelCase_           =jax.nn.softmax(temp_dist_warper_sharper(_lowerCAmelCase		,	scores.copy()		,	cur_len=_lowerCAmelCase							)		,	axis=-1							)
         UpperCAmelCase_           =jax.nn.softmax(temp_dist_warper_smoother(_lowerCAmelCase		,	scores.copy()		,	cur_len=_lowerCAmelCase							)		,	axis=-1							)

         # uniform distribution stays uniform
         self.assertTrue(jnp.allclose(probs[0, :]		,	warped_prob_sharp[0, :]		,	atol=1e-3							)							)
         self.assertTrue(jnp.allclose(probs[0, :]		,	warped_prob_smooth[0, :]		,	atol=1e-3							)							)

         # sharp peaks get higher, valleys get lower
         self.assertLess(probs[1, :].max()		,	warped_prob_sharp[1, :].max()							)
         self.assertGreater(probs[1, :].min()		,	warped_prob_sharp[1, :].min()							)

         # smooth peaks get lower, valleys get higher
         self.assertGreater(probs[1, :].max()		,	warped_prob_smooth[1, :].max()							)
         self.assertLess(probs[1, :].min()		,	warped_prob_smooth[1, :].min()							)



      def 						lowerCAmelCase__     (						self:    Optional[Any]							)    ->     int:
         '''simple docstring'''

         UpperCAmelCase_           =None
         UpperCAmelCase_           =10
         UpperCAmelCase_           =2

         # create ramp distribution
         UpperCAmelCase_           =np.broadcast_to(np.arange(_lowerCAmelCase							)[None, :]		,	(batch_size, vocab_size)							).copy()
         UpperCAmelCase_           =ramp_logits[1:, : vocab_size // 2] + vocab_size

         UpperCAmelCase_           =FlaxTopKLogitsWarper(3							)

         UpperCAmelCase_           =top_k_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # check that correct tokens are filtered
         self.assertListEqual(jnp.isinf(scores[0]							).tolist()		,	7 * [True] + 3 * [False]							)
         self.assertListEqual(jnp.isinf(scores[1]							).tolist()		,	2 * [True] + 3 * [False] + 5 * [True]							)

         # check special case
         UpperCAmelCase_           =5
         UpperCAmelCase_           =FlaxTopKLogitsWarper(top_k=1		,	filter_value=0.0		,	min_tokens_to_keep=3							)

         UpperCAmelCase_           =np.broadcast_to(np.arange(_lowerCAmelCase							)[None, :]		,	(batch_size, length)							).copy()
         UpperCAmelCase_           =top_k_warp_safety_check(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
         self.assertListEqual((scores == 0.0).sum(axis=-1							).tolist()		,	[2, 2]							)



      def 						lowerCAmelCase__     (						self:    Tuple							)    ->     List[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =None
         UpperCAmelCase_           =10
         UpperCAmelCase_           =2

         # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
         UpperCAmelCase_           =np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]							)							)

         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.8							)
         UpperCAmelCase_           =np.exp(top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)							)

         # dist should be filtered to keep min num values so that sum is >= top_p
         # exp (-inf) => 0
         UpperCAmelCase_           =np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]							)
         self.assertTrue(np.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # check edge cases with negative and extreme logits
         UpperCAmelCase_           =np.broadcast_to(np.arange(_lowerCAmelCase							)[None, :]		,	(batch_size, vocab_size)							).copy() - (
             vocab_size // 2
         )

         # make ramp_logits more extreme
         UpperCAmelCase_           =ramp_logits[1] * 1_00.0

         # make sure at least 2 tokens are kept
         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.9		,	min_tokens_to_keep=2		,	filter_value=0.0							)
         UpperCAmelCase_           =top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
         self.assertListEqual((filtered_dist != 0.0).sum(axis=-1							).tolist()		,	[3, 2]							)



      def 						lowerCAmelCase__     (						self:    Any							)    ->     Any:
         '''simple docstring'''

         UpperCAmelCase_           =20
         UpperCAmelCase_           =4
         UpperCAmelCase_           =0

         UpperCAmelCase_           =FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=_lowerCAmelCase							)

         # check that min length is applied at length 5
         UpperCAmelCase_           =ids_tensor((batch_size, 20)		,	vocab_size=20							)
         UpperCAmelCase_           =5
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =min_dist_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,	4 * [-float("inf"							)]							)

         # check that min length is not applied anymore at length 15
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =15
         UpperCAmelCase_           =min_dist_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertFalse(jnp.isinf(_lowerCAmelCase							).any()							)



      def 						lowerCAmelCase__     (						self:    Any							)    ->     Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_           =20
         UpperCAmelCase_           =4
         UpperCAmelCase_           =0

         UpperCAmelCase_           =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase							)

         # check that all scores are -inf except the bos_token_id score
         UpperCAmelCase_           =ids_tensor((batch_size, 1)		,	vocab_size=20							)
         UpperCAmelCase_           =1
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]							).all()							)
         self.assertListEqual(scores[:, bos_token_id].tolist()		,	4 * [0]							)  # score for bos_token_id shold be zero

         # check that bos_token_id is not forced if current length is greater than 1
         UpperCAmelCase_           =3
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertFalse(jnp.isinf(_lowerCAmelCase							).any()							)



      def 						lowerCAmelCase__     (						self:    str							)    ->     Any:
         '''simple docstring'''

         UpperCAmelCase_           =20
         UpperCAmelCase_           =4
         UpperCAmelCase_           =0
         UpperCAmelCase_           =5

         UpperCAmelCase_           =FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCAmelCase		,	eos_token_id=_lowerCAmelCase							)

         # check that all scores are -inf except the eos_token_id when max_length is reached
         UpperCAmelCase_           =ids_tensor((batch_size, 4)		,	vocab_size=20							)
         UpperCAmelCase_           =4
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]							).all()							)
         self.assertListEqual(scores[:, eos_token_id].tolist()		,	4 * [0]							)  # score for eos_token_id should be zero

         # check that eos_token_id is not forced if max_length is not reached
         UpperCAmelCase_           =3
         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =logits_processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         self.assertFalse(jnp.isinf(_lowerCAmelCase							).any()							)



      def 						lowerCAmelCase__     (						self:    List[str]							)    ->     Optional[int]:
         '''simple docstring'''

         UpperCAmelCase_           =4
         UpperCAmelCase_           =10
         UpperCAmelCase_           =15
         UpperCAmelCase_           =2
         UpperCAmelCase_           =1
         UpperCAmelCase_           =15

         # dummy input_ids and scores
         UpperCAmelCase_           =ids_tensor((batch_size, sequence_length)		,	_lowerCAmelCase							)
         UpperCAmelCase_           =input_ids.copy()

         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =scores.copy()

         # instantiate all dist processors
         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=0.5							)
         UpperCAmelCase_           =FlaxTopKLogitsWarper(3							)
         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.8							)

         # instantiate all logits processors
         UpperCAmelCase_           =FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCAmelCase		,	eos_token_id=_lowerCAmelCase							)

         UpperCAmelCase_           =10

         # no processor list
         UpperCAmelCase_           =temp_dist_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =top_k_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =min_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =bos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
         UpperCAmelCase_           =eos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # with processor list
         UpperCAmelCase_           =FlaxLogitsProcessorList(
             [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
         UpperCAmelCase_           =processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)

         # scores should be equal
         self.assertTrue(jnp.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()		,	input_ids_comp.tolist()							)






      def 						lowerCAmelCase__     (						self:    str							)    ->     Tuple:
         '''simple docstring'''

         UpperCAmelCase_           =4
         UpperCAmelCase_           =10
         UpperCAmelCase_           =15
         UpperCAmelCase_           =2
         UpperCAmelCase_           =1
         UpperCAmelCase_           =15

         # dummy input_ids and scores
         UpperCAmelCase_           =ids_tensor((batch_size, sequence_length)		,	_lowerCAmelCase							)
         UpperCAmelCase_           =input_ids.copy()

         UpperCAmelCase_           =self._get_uniform_logits(_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =scores.copy()

         # instantiate all dist processors
         UpperCAmelCase_           =FlaxTemperatureLogitsWarper(temperature=0.5							)
         UpperCAmelCase_           =FlaxTopKLogitsWarper(3							)
         UpperCAmelCase_           =FlaxTopPLogitsWarper(0.8							)

         # instantiate all logits processors
         UpperCAmelCase_           =FlaxMinLengthLogitsProcessor(min_length=10		,	eos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase							)
         UpperCAmelCase_           =FlaxForcedEOSTokenLogitsProcessor(max_length=_lowerCAmelCase		,	eos_token_id=_lowerCAmelCase							)

         UpperCAmelCase_           =10

         # no processor list
         def run_no_processor_list(_lowerCAmelCase:    int		,	_lowerCAmelCase:    Optional[Any]		,	_lowerCAmelCase:    Dict							):
            UpperCAmelCase_           =temp_dist_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =top_k_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =top_p_warp(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =min_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =bos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            UpperCAmelCase_           =eos_dist_proc(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            return scores

         # with processor list
         def run_processor_list(_lowerCAmelCase:    Optional[int]		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    Optional[int]							):
            UpperCAmelCase_           =FlaxLogitsProcessorList(
                [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]							)
            UpperCAmelCase_           =processor(_lowerCAmelCase		,	_lowerCAmelCase		,	cur_len=_lowerCAmelCase							)
            return scores

         UpperCAmelCase_           =jax.jit(_lowerCAmelCase							)
         UpperCAmelCase_           =jax.jit(_lowerCAmelCase							)

         UpperCAmelCase_           =jitted_run_no_processor_list(_lowerCAmelCase		,	_lowerCAmelCase		,	_lowerCAmelCase							)
         UpperCAmelCase_           =jitted_run_processor_list(_lowerCAmelCase		,	_lowerCAmelCase		,	_lowerCAmelCase							)

         # scores should be equal
         self.assertTrue(jnp.allclose(_lowerCAmelCase		,	_lowerCAmelCase		,	atol=1e-3							)							)

         # input_ids should never be changed
         self.assertListEqual(input_ids.tolist()		,	input_ids_comp.tolist()							)








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


SCREAMING_SNAKE_CASE							    =							logging.getLogger(__name__)

SCREAMING_SNAKE_CASE							    =							5_0  # max width of layer names
SCREAMING_SNAKE_CASE							    =							7_0  # max width of quantizer names




def  a   (lowerCAmelCase__   ):
   __a					=  parser.add_argument_group("""quant_trainer arguments"""   )
   group.add_argument("""--wprec"""		,				type=lowerCAmelCase__		,				default=8		,				help="""weight precision"""   )
   group.add_argument("""--aprec"""		,				type=lowerCAmelCase__		,				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=lowerCAmelCase__		,				nargs="""+"""		,				help="""disable quantizers by keyword"""   )
   group.add_argument("""--quant-disable-layer-module"""		,				type=lowerCAmelCase__		,				help="""disable quantizers by keyword under layer."""   )
   group.add_argument("""--quant-enable-layer-module"""		,				type=lowerCAmelCase__		,				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=lowerCAmelCase__		,				type=lowerCAmelCase__		,				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=lowerCAmelCase__		,				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  a   (lowerCAmelCase__   ):

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

   __a					=  QuantDescriptor(num_bits=args.aprec		,				calib_method=lowerCAmelCase__   )
   __a					=  QuantDescriptor(num_bits=args.wprec		,				axis=(None if args.quant_per_tensor else (0,))   )
   quant_nn.QuantLinear.set_default_quant_desc_input(lowerCAmelCase__   )
   quant_nn.QuantLinear.set_default_quant_desc_weight(lowerCAmelCase__   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__=False		,				lowerCAmelCase__=False   ):
   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(lowerCAmelCase__		,				["""embeddings"""]		,				which="""weight"""		,				_disabled=lowerCAmelCase__   )

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

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

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

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

      if args.recalibrate_weights:
         recalibrate_weights(lowerCAmelCase__   )

      if args.fuse_qkv:
         fuse_qkv(lowerCAmelCase__		,				lowerCAmelCase__   )

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

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




def  a   (lowerCAmelCase__   ):
   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  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):
   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(lowerCAmelCase__   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):

   def fusea(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   ):
      for mod in [qq, qk, qv]:
         if not hasattr(lowerCAmelCase__		,				"""_amax"""   ):
            print("""          WARNING: NO AMAX BUFFER"""   )
            return
      __a					=  qq._amax.detach().item()
      __a					=  qk._amax.detach().item()
      __a					=  qv._amax.detach().item()

      __a					=  max(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   )
      qq._amax.fill_(lowerCAmelCase__   )
      qk._amax.fill_(lowerCAmelCase__   )
      qv._amax.fill_(lowerCAmelCase__   )
      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  a   (lowerCAmelCase__		,				lowerCAmelCase__   ):

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




def  a   (lowerCAmelCase__   ):

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




def  a   (lowerCAmelCase__   ):

   for name, mod in model.named_modules():
      if hasattr(lowerCAmelCase__		,				"""_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)
         __a					=  set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis   )
         __a					=  set(range(len(mod.weight.size()   )   )   ) - axis_set
         __a					=  pytorch_quantization.utils.reduce_amax(mod.weight		,				axis=lowerCAmelCase__		,				keepdims=lowerCAmelCase__   ).detach()
         logger.info(f'''RECALIB: {name:{name_width}} {mod._weight_quantizer._amax.flatten()} -> {amax.flatten()}'''   )
         __a					=  amax




def  a   (lowerCAmelCase__		,				lowerCAmelCase__=25		,				lowerCAmelCase__=180		,				lowerCAmelCase__=None   ):

   if ignore is None:
      __a					=  []
   elif not isinstance(lowerCAmelCase__		,				lowerCAmelCase__   ):
      __a					=  [ignore]

   __a					=  0
   for name, mod in model.named_modules():
      if not hasattr(lowerCAmelCase__		,				"""weight"""   ):
         continue
      __a					=  max(lowerCAmelCase__		,				len(lowerCAmelCase__   )   )

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




def  a   (lowerCAmelCase__   ):
   __a					=  0
   for name, mod in model.named_modules():
      if isinstance(lowerCAmelCase__		,				pytorch_quantization.nn.TensorQuantizer   ):
         print(f'''{name:80} {mod}'''   )
         count += 1
   print(f'''{count} TensorQuantizers found in model'''   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   ):
   __a					=  getattr(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   )
   if quantizer_mod is not None:
      assert hasattr(lowerCAmelCase__		,				lowerCAmelCase__   )
      setattr(lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__   )
   else:
      logger.warning(f'''{name} has no {quantizer}'''   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__="both"		,				**lowerCAmelCase__   ):
   __a					=  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(lowerCAmelCase__		,				lowerCAmelCase__		,				"""_input_quantizer"""		,				lowerCAmelCase__		,				lowerCAmelCase__   )
      if which in ["weight", "both"]:
         set_quantizer(lowerCAmelCase__		,				lowerCAmelCase__		,				"""_weight_quantizer"""		,				lowerCAmelCase__		,				lowerCAmelCase__   )
   logger.info(lowerCAmelCase__   )




def  a   (lowerCAmelCase__		,				lowerCAmelCase__		,				**lowerCAmelCase__   ):

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



