






"""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


lowerCamelCase      =		logging.getLogger(__name__)

lowerCamelCase      =		50  # max width of layer names
lowerCamelCase      =		70  # max width of quantizer names
def      a__			(   lowerCAmelCase__   ):
      UpperCAmelCase_    				=							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":
            UpperCAmelCase_    				=							"max"
      elif args.calibrator == "percentile":
            if args.percentile is None:
                  raise ValueError("Specify --percentile when using percentile calibrator"   )
            UpperCAmelCase_    				=							"histogram"
      elif args.calibrator == "mse":
            UpperCAmelCase_    				=							"histogram"
      else:
            raise ValueError(f"""Invalid calibrator {args.calibrator}"""   )

      UpperCAmelCase_    				=							QuantDescriptor(num_bits=args.aprec			,    calib_method=lowerCAmelCase__   )
      UpperCAmelCase_    				=							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
            UpperCAmelCase_    				=							qq._amax.detach().item()
            UpperCAmelCase_    				=							qk._amax.detach().item()
            UpperCAmelCase_    				=							qv._amax.detach().item()

            UpperCAmelCase_    				=							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"   ):
                  UpperCAmelCase_    				=							mod._input_quantizer._amax.data.detach().item()
                  mod._input_quantizer._amax.data.detach().clamp_(max=lowerCAmelCase__   )
                  UpperCAmelCase_    				=							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:
                  UpperCAmelCase_    				=							mod.weight.shape[0]
                  UpperCAmelCase_    				=							mod._weight_quantizer._amax.detach()
                  UpperCAmelCase_    				=							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)
                  UpperCAmelCase_    				=							set() if mod._weight_quantizer.axis is None else set(mod._weight_quantizer.axis   )
                  UpperCAmelCase_    				=							set(range(len(mod.weight.size()   )   )   ) - axis_set
                  UpperCAmelCase_    				=							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()}"""   )
                  UpperCAmelCase_    				=							amax
def      a__			(   lowerCAmelCase__			,    lowerCAmelCase__=25			,    lowerCAmelCase__=180			,    lowerCAmelCase__=None   ):

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

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

      for name, mod in model.named_modules():
            UpperCAmelCase_    				=							getattr(lowerCAmelCase__			,    "_input_quantizer"			,    lowerCAmelCase__   )
            UpperCAmelCase_    				=							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
            UpperCAmelCase_    				=							f"""Act:{input_q.extra_repr()}"""
            UpperCAmelCase_    				=							f"""Wgt:{weight_q.extra_repr()}"""
            UpperCAmelCase_    				=							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__   ):
      UpperCAmelCase_    				=							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__   ):
      UpperCAmelCase_    				=							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__   ):
      UpperCAmelCase_    				=							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__   ):
                              UpperCAmelCase_    				=							f"""Warning: changing {name:{name_width}}"""
                              for k, v in kwargs.items():
                                    s += f""" {k}={v}"""
                                    setattr(lowerCAmelCase__			,    lowerCAmelCase__			,    lowerCAmelCase__   )
                              logger.info(lowerCAmelCase__   )









"""simple docstring"""





import colorsys

from PIL import Image  # type: ignore
def      a__			(   lowerCAmelCase__			,    lowerCAmelCase__			,    lowerCAmelCase__   ):
      UpperCAmelCase_    				=							x
      UpperCAmelCase_    				=							y
      for step in range(lowerCAmelCase__   ):  # noqa: B007
            UpperCAmelCase_    				=							a * a - b * b + x
            UpperCAmelCase_    				=							2 * a * b + y
            UpperCAmelCase_    				=							a_new

            # divergence happens for all complex number with an absolute value
            # greater than 4
            if a * a + b * b > 4:
                  break
      return step / (max_step - 1)
def      a__			(   lowerCAmelCase__   ):
      if distance == 1:
            return (0, 0, 0)
      else:
            return (255, 255, 255)
def      a__			(   lowerCAmelCase__   ):
      if distance == 1:
            return (0, 0, 0)
      else:
            return tuple(round(i * 255   ) for i in colorsys.hsv_to_rgb(lowerCAmelCase__			,    1			,    1   )   )
def      a__			(   lowerCAmelCase__ = 800			,    lowerCAmelCase__ = 600			,    lowerCAmelCase__ = -0.6			,    lowerCAmelCase__ = 0			,    lowerCAmelCase__ = 3.2			,    lowerCAmelCase__ = 50			,    lowerCAmelCase__ = True			,    ):
      UpperCAmelCase_    				=							Image.new("RGB"			,    (image_width, image_height)   )
      UpperCAmelCase_    				=							img.load()

      # loop through the image-coordinates
      for image_x in range(lowerCAmelCase__   ):
            for image_y in range(lowerCAmelCase__   ):
                  # determine the figure-coordinates based on the image-coordinates
                  UpperCAmelCase_    				=							figure_width / image_width * image_height
                  UpperCAmelCase_    				=							figure_center_x + (image_x / image_width - 0.5) * figure_width
                  UpperCAmelCase_    				=							figure_center_y + (image_y / image_height - 0.5) * figure_height

                  UpperCAmelCase_    				=							get_distance(lowerCAmelCase__			,    lowerCAmelCase__			,    lowerCAmelCase__   )

                  # color the corresponding pixel based on the selected coloring-function
                  if use_distance_color_coding:
                        UpperCAmelCase_    				=							get_color_coded_rgb(lowerCAmelCase__   )
                  else:
                        UpperCAmelCase_    				=							get_black_and_white_rgb(lowerCAmelCase__   )

      return img


if __name__ == "__main__":
  import doctest

  doctest.testmod()

  # colored version, full figure
  lowerCamelCase      =		get_image()

  # uncomment for colored version, different section, zoomed in
  # img = get_image(figure_center_x = -0.6, figure_center_y = -0.4,
  # figure_width = 0.8)

  # uncomment for black and white version, full figure
  # img = get_image(use_distance_color_coding = False)

  # uncomment to save the image
  # img.save("mandelbrot.png")

  img.show()




