



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    _snake_case (       __snake_case    )     ->     List[Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	Dict				      =		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    _snake_case (       __snake_case    )     ->     Optional[int]:






      '''simple docstring'''







      if args.calibrator == "max":
            UpperCAmelCase_  :	List[Any]				      =		"max"
      elif args.calibrator == "percentile":
            if args.percentile is None:
                  raise ValueError("Specify --percentile when using percentile calibrator"    )
            UpperCAmelCase_  :	List[Any]				      =		"histogram"
      elif args.calibrator == "mse":
            UpperCAmelCase_  :	Union[str, Any]				      =		"histogram"
      else:
            raise ValueError(F"""Invalid calibrator {args.calibrator}"""    )

      UpperCAmelCase_  :	List[Any]				      =		QuantDescriptor(num_bits=args.aprec					,  calib_method=__snake_case    )
      UpperCAmelCase_  :	Any				      =		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    _snake_case (       __snake_case					,  __snake_case					,  __snake_case=False					,  __snake_case=False    )     ->     str:






      '''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    _snake_case (       __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    _snake_case (       __snake_case					,  __snake_case    )     ->     Union[str, 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(__snake_case    )







def    _snake_case (       __snake_case					,  __snake_case    )     ->     Any:






      '''simple docstring'''







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

            UpperCAmelCase_  :	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    _snake_case (       __snake_case					,  __snake_case    )     ->     Any:






      '''simple docstring'''







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







def    _snake_case (       __snake_case    )     ->     Optional[int]:






      '''simple docstring'''







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







def    _snake_case (       __snake_case    )     ->     int:






      '''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)
                  UpperCAmelCase_  :	Union[str, Any]				      =		set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis    )
                  UpperCAmelCase_  :	Optional[int]				      =		set(range(len(mod.weight.size()    )    )    ) - axis_set
                  UpperCAmelCase_  :	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()}"""    )
                  UpperCAmelCase_  :	Union[str, Any]				      =		amax







def    _snake_case (       __snake_case					,  __snake_case=2_5					,  __snake_case=1_8_0					,  __snake_case=None    )     ->     Union[str, Any]:






      '''simple docstring'''







      if ignore is None:
            UpperCAmelCase_  :	Tuple				      =		[]
      elif not isinstance(__snake_case					,  __snake_case    ):
            UpperCAmelCase_  :	str				      =		[ignore]

      UpperCAmelCase_  :	Any				      =		0
      for name, mod in model.named_modules():
            if not hasattr(__snake_case					,  "weight"    ):
                  continue
            UpperCAmelCase_  :	Optional[Any]				      =		max(__snake_case					,  len(__snake_case    )    )

      for name, mod in model.named_modules():
            UpperCAmelCase_  :	List[Any]				      =		getattr(__snake_case					,  "_input_quantizer"					,  __snake_case    )
            UpperCAmelCase_  :	List[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
            UpperCAmelCase_  :	Optional[Any]				      =		F"""Act:{input_q.extra_repr()}"""
            UpperCAmelCase_  :	Optional[Any]				      =		F"""Wgt:{weight_q.extra_repr()}"""
            UpperCAmelCase_  :	Optional[Any]				      =		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    _snake_case (       __snake_case    )     ->     List[Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	Dict				      =		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    _snake_case (       __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case    )     ->     Optional[Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	Optional[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    _snake_case (       __snake_case					,  __snake_case					,  __snake_case="both"					,  **__snake_case    )     ->     Any:






      '''simple docstring'''







      UpperCAmelCase_  :	Union[str, Any]				      =		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    _snake_case (       __snake_case					,  __snake_case					,  **__snake_case    )     ->     List[str]:






      '''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    ):
                              UpperCAmelCase_  :	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 __future__ import annotations

from collections import deque
from collections.abc import Iterator
from dataclasses import dataclass


@dataclass
class   snake_case_ :




       """simple docstring"""






       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				42


class   snake_case_ :




       """simple docstring"""






       def __init__(     self     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	list[list[Edge]]				      =		[[] for _ in range(lowercase)]
             UpperCAmelCase_  :	str				      =		size



       def __getitem__(     self     ,lowercase):

             """simple docstring"""





             return iter(self._graph[vertex])



       @property
       def  A_   (     self):

             """simple docstring"""





             return self._size



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             if weight not in (0, 1):
                   raise ValueError("Edge weight must be either 0 or 1.")

             if to_vertex < 0 or to_vertex >= self.size:
                   raise ValueError("Vertex indexes must be in [0; size).")

             self._graph[from_vertex].append(Edge(lowercase     ,lowercase))



       def  A_   (     self     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		deque([start_vertex])
             UpperCAmelCase_  :	list[int | None]				      =		[None] * self.size
             UpperCAmelCase_  :	Dict				      =		0

             while queue:
                   UpperCAmelCase_  :	List[Any]				      =		queue.popleft()
                   UpperCAmelCase_  :	Any				      =		distances[current_vertex]
                   if current_distance is None:
                         continue

                   for edge in self[current_vertex]:
                         UpperCAmelCase_  :	Tuple				      =		current_distance + edge.weight
                         UpperCAmelCase_  :	Optional[Any]				      =		distances[edge.destination_vertex]
                         if (
                             isinstance(lowercase     ,lowercase)
                             and new_distance >= dest_vertex_distance
                         ):
                               continue
                         UpperCAmelCase_  :	Tuple				      =		new_distance
                         if edge.weight == 0:
                               queue.appendleft(edge.destination_vertex)
                         else:
                               queue.append(edge.destination_vertex)

             if distances[finish_vertex] is None:
                   raise ValueError("No path from start_vertex to finish_vertex.")

             return distances[finish_vertex]


if __name__ == "__main__":
   import doctest

   doctest.testmod()
