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 _load_from_state_dict(
      self,
      state_dict,
      prefix,
      local_metadata,
      strict,
      missing_keys,
      unexpected_keys,
      error_msgs,
  ):
      self.cluster_dtype = local_metadata["cluster_dtype"]
      state_dict_buffers_are_placeholders = local_metadata[
          "buffers_are_placeholders"
      ]
  
      if (
          not self.buffers_are_placeholders
          and state_dict_buffers_are_placeholders
      ):
          raise ValueError(
              f"Trying to reload an uninitialized state dict onto an initialized module: {prefix[:-1]}"
          )
  
      if (
          self.buffers_are_placeholders
          and not state_dict_buffers_are_placeholders
      ):
          # We only change the size of the placeholders if we intend to reload a proper checkpoint
          # onto an uninitialized module. In the other cases, we expect the state dict and the module to be compatible.
          self.centroids = _torch.empty(
              state_dict[prefix + "centroids"].size(),
              device=self.centroids.device,
          )
          self.labels = _torch.empty(
              state_dict[prefix + "labels"].size(),
              device=self.labels.device,
          )
          self.fake_palett_enabled = _torch.empty(
              state_dict[prefix + "fake_palett_enabled"].size(),
              device=self.labels.device,
          )
  
      self.buffers_are_placeholders = (
          state_dict_buffers_are_placeholders
      )
  
      _Partitioner._load_from_state_dict_(
          self,
          state_dict,
          prefix + "palett.",
          local_metadata,
          strict,
          missing_keys,
          unexpected_keys,
          error_msgs,
      )
  
      if self.need_to_quantize:
          # We will go through FakeQuantize._load_from_state_dict and then nn.Module._load_from_state_dict
          super()._load_from_state_dict(
              state_dict,
              prefix,
              local_metadata,
              strict,
              missing_keys,
              unexpected_keys,
              error_msgs,
          )
      else:
          # Jump FakeQuantize and go to nn.Module directly
          super(_FakeQuantize, self)._load_from_state_dict(
              state_dict,
              prefix,
              local_metadata,
              strict,
              missing_keys,
              unexpected_keys,
              error_msgs,
          )
  
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 7 for calling the FakePalettize obj's underlying member method code with example usage: obj._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) without return types