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 forward(self, weights: _torch.Tensor):
      if self.partition_size == 0:
          forwarded_weights = super().forward(weights)
          if self.fake_palett_enabled[0] == 1:
              with _torch.no_grad():
                  (
                      quant_centroids,
                      quant_labels,
                  ) = forwarded_weights.unique(return_inverse=True)
                  self.centroids = _torch.stack(
                      [quant_centroids.view(-1, self.cluster_dim)]
                  )
                  self.labels = _torch.stack([quant_labels])
      else:
          forwarded_weights = weights.clone()
  
          if self.fake_palett_enabled[0] == 1:
              if not self.partitions:
                  self.init_partitions(weights.detach())
                  self.centroids = _torch.stack(self.centroids_init)
                  self.labels = _torch.stack(self.labels_init)
                  self.buffers_are_placeholders = False
  
              if self.training:
                  forwarded_weights = self.diff_palettize(
                      forwarded_weights
                  )
              else:
                  forwarded_weights = self.palettize(forwarded_weights)
          else:
              forwarded_weights = super().forward(weights)
  
      if self.cluster_dtype == "f16":
          forwarded_weights = forwarded_weights.to(_torch.float16).to(
              weights.dtype
          )
      elif self.cluster_dtype == "b16":
          forwarded_weights = forwarded_weights.to(_torch.bfloat16).to(
              weights.dtype
          )
  
      return forwarded_weights
  
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 6 for calling the FakePalettize obj's underlying member method code with example usage: obj.forward(weights) and returns: forwarded_weights