class_docstr: A class that implements palettization algorithm described in
  `DKM: Differentiable K-Means Clustering Layer for Neural Network Compression
  <https://arxiv.org/abs/2108.12659>`_. It clusters the weights
  using a differentiable version of ``k-means``, allowing the look-up-table (LUT)
  and indices of palettized weights to be learnt using a gradient-based optimization
  algorithm such as SGD.
  
  Extends :py:class:`torch.quantization.FakeQuantize` to add support for
  palettization.
  
  Example:
          .. code-block:: python
  
              from collections import OrderedDict
              import torch
              import torch.nn as nn
              import coremltools.optimize.torch.palettization as palett
  
              model = nn.Sequential(
                  OrderedDict(
                      [
                          ("linear1", nn.Linear(4, 5)),
                          ("sigmoid1", nn.Sigmoid()),
                          ("linear2", nn.Linear(5, 4)),
                          ("sigmoid2", nn.Sigmoid),
                      ]
                  )
              )
  
              fq_activation = nn.Identity
              fq_weight = palett.FakePalettize.with_args(
                  observer=torch.quantization.MovingAveragePerChannelMinMaxObserver.with_args(
                      quant_min=-128, quant_max=127, dtype=torch.qint8
                  ),
                  n_bits=2,
                  cluster_dim=1,
              )
              model.linear2.qconfig = torch.quantization.QConfig(
                  activation=fq_activation, weight=fq_weight
              )
  
              palettized_model = palett.prepare_palettizer(model)
  
              train_model(palettized_model)
  
              palettized_converted_model = palett.finalize(palettized_model)
  
  
  Args:
      observer (:obj:`torch.ao.quantization.observer.ObserverBase`): Observer for quantizing the ``LUT``.
      n_bits (:obj:`int`): Number of palettization bits. There would be :math:`2^{n\_bits}` unique weights in the ``LUT``.
      cluster_dim (:obj:`int`): Dimensionality of centroids to use for clustering.
      quant_min (:obj:`int`): The minimum allowable quantized value.
      quant_max (:obj:`int`): The maximum allowable quantized value.
      cluster_dtype (:obj:`str`): String that decides whether to quantize the ``LUT`` or not. The following are the ``str``
          LUT quantization combinations: (``u8``, ``uint8``), (``i8``, ``int8``), and (``f16``, ``float16``).
      advanced_options (:obj:`dict`): Advanced options to configure the palettization algorithm.
      observer_kwargs (optional): Arguments for the observer module.
  
  .. note::
      Allowed keys for ``advanced_options`` are the parameters listed as ``optional`` in
      :py:class:`ModuleDKMPalettizerConfig`, besides the ones already covered by other parameters in this class.
code: def __init__(
      self,
      observer: _ObserverBase,
      n_bits: int,
      cluster_dim: int,
      quant_min: int = -128,
      quant_max: int = 127,
      cluster_dtype: str = "f32",
      advanced_options: dict = {},
      **observer_kwargs,
  ):
      partition_size = advanced_options.get(
          "partition_size",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["partition_size"],
      )
      cluster_permute = advanced_options.get(
          "cluster_permute",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["cluster_permute"],
      )
      palett_max_mem = advanced_options.get(
          "palett_max_mem",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["palett_max_mem"],
      )
      kmeans_max_iter = advanced_options.get(
          "kmeans_max_iter",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["kmeans_max_iter"],
      )
      prune_threshold = advanced_options.get(
          "prune_threshold",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["prune_threshold"],
      )
      kmeans_init = advanced_options.get(
          "kmeans_init",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["kmeans_init"],
      )
      kmeans_opt1d_threshold = advanced_options.get(
          "kmeans_opt1d_threshold",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS[
              "kmeans_opt1d_threshold"
          ],
      )
      enforce_zero = advanced_options.get(
          "enforce_zero",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["enforce_zero"],
      )
      palett_mode = advanced_options.get(
          "palett_mode",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["palett_mode"],
      )
      palett_cluster_tol = advanced_options.get(
          "palett_cluster_tol",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["palett_cluster_tol"],
      )
      palett_tau = advanced_options.get(
          "palett_tau",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["palett_tau"],
      )
      palett_epsilon = advanced_options.get(
          "palett_epsilon",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["palett_epsilon"],
      )
      palett_lambda = advanced_options.get(
          "palett_lambda",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["palett_lambda"],
      )
      add_extra_centroid = advanced_options.get(
          "add_extra_centroid",
          DEFAULT_PALETTIZATION_ADVANCED_OPTIONS["add_extra_centroid"],
      )
  
      self._target_module_level_sparsity = 0.0
  
      _FakeQuantize.__init__(
          self, observer, quant_min, quant_max, **observer_kwargs
      )
      _Partitioner.__init__(
          self,
          n_bits,
          enforce_zero,
          partition_size,
          cluster_dim,
          cluster_permute,
          palett_tau,
          kmeans_init,
          prune_threshold,
          kmeans_opt1d_threshold,
          add_extra_centroid,
      )
  
      self.cluster_dtype = cluster_dtype
      self.add_extra_centroid = add_extra_centroid
      self.need_to_quantize = self.cluster_dtype in ["i8", "u8", "f16"]
      self.autograd_graph = (
          hasattr(_torch.autograd, "graph") and palett_max_mem < 1.0
      )
      self.palett_max_mem = palett_max_mem
      self.palett_cluster_tol = palett_cluster_tol
      self.kmeans_max_iter = kmeans_max_iter
      self.palett_mode = palett_mode
      self.palett_tau = palett_tau
      self.palett_epsilon = palett_epsilon
      self.palett_lambda = palett_lambda
      self.n_bits = n_bits
      self.cluster_dim = cluster_dim
      self.kmeans_init = kmeans_init
      # Temporary create placeholder buffers that will get replaced with proper centroids on the first forward,
      # or when we reload a checkpoint. Having placeholder values is useful to maintain the structure of the state
      # dict constant.
      self.register_buffer("centroids", _torch.rand([1]))
      self.register_buffer("labels", _torch.rand([1]))
      # During init, we would want the fake_palett_enabled flag to be False, i.e. to be at a state of 0. Also, we
      # would have set the fake_quant_enabled and observer_enabled to be 0 as well so that palettizer does nothing
      # until the first milestone.
      self.register_buffer(
          "fake_palett_enabled", _torch.tensor([0], dtype=_torch.uint8)
      )
      self.disable_fake_quant()
      self.disable_observer()
      self.buffers_are_placeholders = True
  
imports: import contextlib
  import gc
  import torch
  import torch.nn.functional
  from torch.ao.quantization.observer import ObserverBase
  from torch.quantization import FakeQuantize
  from ._efficient_kmeans import _EfficientKMeans
  from ._fake_palettizer_tensor_hook import _FakePalettizationTensorHook
  from ._partitioner import _Partitioner
  from .palettization_config import DEFAULT_PALETTIZATION_ADVANCED_OPTIONS
label_desc: Use image node_id 1 to create a new FakePalettize object from inherited base classes: _FakeQuantize, _Partitioner with example: obj = FakePalettize(observer, n_bits, cluster_dim, quant_min, quant_max, cluster_dtype, advanced_options)