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.
docstr: Method called to run the differentiable k-means operation.
docstr_tok: ["Method","called","to","run","the","differentiable","k-means","operation","."]
code: def diff_palettize(self, weights: _torch.Tensor):
      """
      Method called to run the differentiable k-means operation.
      """
      use_cpu_if_cuda_available = False
      if _torch.cuda.is_available():
          t = _torch.cuda.get_device_properties(
              weights.device
          ).total_memory
          a = _torch.cuda.memory_allocated(weights.device)
          use_cpu_if_cuda_available = (
              a / t
          ) > self.palett_max_mem and self.autograd_graph
          if use_cpu_if_cuda_available:
              if _FakePalettizationTensorHook.gc_trigger is None:
                  _FakePalettizationTensorHook.gc_trigger = True
  
      if _FakePalettizationTensorHook.gc_trigger:
          gc.collect()
  
      auto_grad_graph_on_cpu = (
          _torch.autograd.graph.save_on_cpu(pin_memory=True)
          if use_cpu_if_cuda_available
          else contextlib.nullcontext()
      )
  
      for i, partition in enumerate(self.partitions):
          current_partition_clone = weights[
              partition[0] : partition[1]
          ].clone()
          cX, pad = self.flatten(current_partition_clone)
  
          with _torch.no_grad():
              palett_table = _torch.unique(self.centroids[i], dim=0)
              if (
                  len(palett_table)
                  < self.n_clusters[i] * self.palett_cluster_tol
              ):
                  # We use n_init as 3 so as to not spend a lot of time running this operation
                  kmeans = _EfficientKMeans(
                      n_clusters=self.n_clusters[i],
                      init="kmeans++",
                      labels=self.labels[i],
                      n_init=3,
                      max_iter=1,
                  )
                  kmeans.kmeans_pp(3, cX, 0)
                  self.centroids[i] = kmeans.cluster_centers_
  
          centroids = self.centroids[i].clone()
  
          assert not centroids.requires_grad
          last_inertia = None
  
          for j in range(self.kmeans_max_iter):
              if self.autograd_graph:
                  tensor_hook = _FakePalettizationTensorHook(
                      [
                          _torch.Size(
                              [cX.size()[0], centroids.size()[0]]
                          )
                      ],
                      use_cpu_if_cuda_available,
                      f"FakePalettizationTensorHook.{i}.{j}",
                      self.palett_tau,
                  )
                  auto_grad_graph_hook_init = (
                      _torch.autograd.graph.saved_tensors_hooks(
                          tensor_hook.init_pack, tensor_hook.init_unpack
                      )
                  )
                  auto_grad_graph_hook_reuse = (
                      _torch.autograd.graph.saved_tensors_hooks(
                          tensor_hook.reuse_pack,
                          tensor_hook.reuse_unpack,
                      )
                  )
              else:
                  auto_grad_graph_hook_init = contextlib.nullcontext()
                  auto_grad_graph_hook_reuse = contextlib.nullcontext()
  
              with auto_grad_graph_hook_init:
                  x_c_dist = _EfficientKMeans.x_c_dist(cX, centroids)
                  min_error, _ = x_c_dist.min(dim=-1)
  
              with auto_grad_graph_hook_reuse:
                  if "dkm" in self.palett_mode:
                      attention = _F.softmax(
                          -x_c_dist / self.palett_tau, dim=1
                      )
                  elif "gsm" in self.palett_mode:
                      attention = _F.gumbel_softmax(
                          -x_c_dist / self.palett_tau, dim=1
                      )
                  elif "hard" in self.palett_mode:
                      col_idx = x_c_dist.min(dim=1).indices
                      row_idx = _torch.arange(
                          start=0, end=len(col_idx), dtype=_torch.int32
                      ).to(cX.device)
                      attention = _torch.sparse_coo_tensor(
                          _torch.vstack([row_idx, col_idx]),
                          _torch.ones_like(row_idx).to(cX.device),
                          x_c_dist.size(),
                          dtype=x_c_dist.dtype,
                          requires_grad=True,
                      ).to_dense()
  
              assert attention.requires_grad
              attention_sum = attention.sum(dim=0).view(-1, 1)
              attention_sum[attention_sum == 0] = 1e-6
  
              with auto_grad_graph_hook_reuse:
                  centroids = (
                      _torch.matmul(cX.T, attention).T / attention_sum
                  )
  
              with auto_grad_graph_on_cpu:
                  if self.need_to_quantize:
                      centroids = super().forward(centroids)
  
                  assert centroids.requires_grad
  
                  if self.prune_threshold > 0:
                      centroids = _torch.nn.Hardshrink(
                          self.prune_threshold.item()
                      )(centroids)
  
                  if self.enforce_zero[i]:
                      zero_point = (
                          _torch.zeros(centroids[0].size())
                          .to(centroids.device)
                          .unsqueeze(0)
                      )
                      zero_idx = _torch.argmin(
                          _torch.cdist(centroids, zero_point)
                      )
                      centroids[zero_idx] = zero_point
  
              cur_inertia = min_error.sum()
  
              if (
                  last_inertia
                  and abs(last_inertia - cur_inertia)
                  <= self.palett_epsilon
              ):
                  break
  
              last_inertia = cur_inertia
  
          with auto_grad_graph_hook_reuse:
              weights[partition[0] : partition[1]] = self.deflatten(
                  _torch.matmul(attention, centroids),
                  current_partition_clone.size(),
                  pad,
              )
  
              self.centroids[i] = (
                  self.palett_lambda * self.centroids[i]
                  + (1 - self.palett_lambda) * centroids
              ).detach()
              self.labels[i] = attention.detach().max(dim=1)[1].data
  
      return 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 4 for calling the FakePalettize obj's underlying member method code with example usage: obj.diff_palettize(weights) and returns: weights