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 _save_to_state_dict(self, destination, prefix, keep_vars):
      if self.need_to_quantize:
          # Use normal inheritance, go through FakeQuantize._save_to_state_dict
          super()._save_to_state_dict(destination, prefix, keep_vars)
          self.centroids = super().forward(self.centroids)
      else:
          # Skip FakeQuantize._save_to_state_dict and go directly to nn.Module._save_to_state_dict
          super(_FakeQuantize, self)._save_to_state_dict(
              destination, prefix, keep_vars
          )
  
      # State dicts can only contain tensors (for DDP), so store infos in the metadata dict (in particular str)
      destination._metadata[prefix[:-1]][
          "cluster_dtype"
      ] = self.cluster_dtype
      destination._metadata[prefix[:-1]][
          "buffers_are_placeholders"
      ] = self.buffers_are_placeholders
      _Partitioner._save_to_state_dict_(
          self, destination, prefix + "palett.", keep_vars
      )
  
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 8 for calling the FakePalettize obj's underlying member method code with example usage: obj._save_to_state_dict(destination, prefix, keep_vars) without return types