class_docstr: A palettization algorithm based on `"DKM: Differentiable K-Means Clustering Layer for Neural Network
  Compression" <https://arxiv.org/pdf/2108.12659.pdf>`_. 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.
  
  Example:
  
      .. code-block:: python
  
          import torch
          from coremltools.optimize.torch.palettization import (
              DKMPalettizer,
              DKMPalettizerConfig,
              ModuleDKMPalettizerConfig,
          )
  
          # code that defines the pytorch model, loss and optimizer.
          model, loss_fn, optimizer = create_model_loss_and_optimizer()
  
          # initialize the palettizer
          config = DKMPalettizerConfig(global_config=ModuleDKMPalettizerConfig(n_bits=4))
  
          palettizer = DKMPalettizer(model, config)
  
          # prepare the model to insert FakePalettize layers for palettization
          model = palettizer.prepare(inplace=True)
  
          # use palettizer in your PyTorch training loop
          for inputs, labels in data:
              output = model(inputs)
              loss = loss_fn(output, labels)
              loss.backward()
              optimizer.step()
              palettizer.step()
  
          # fold LUT and indices into weights
          model = palettizer.finalize(inplace=True)
  
  Args:
      model (:py:class:`torch.nn.Module`): Model on which the palettizer will act.
      config (:py:class:`DKMPalettizerConfig`): Config which specifies how
          different submodules in the model will be configured for palettization.
          Default config is used when passed as ``None``.
class_docstr_tok: ["A","palettization","algorithm","based","on","`","``","DKM",":","Differentiable","K-Means","Clustering","Layer","for","Neural","Network","Compression","''","<","https",":","\/\/arxiv.org\/pdf\/2108.12659.pdf",">","`","_",".","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",".","Example",":","..","code-block",":",":","python","import","torch","from","coremltools.optimize.torch.palettization","import","(","DKMPalettizer",",","DKMPalettizerConfig",",","ModuleDKMPalettizerConfig",",",")","#","code","that","defines","the","pytorch","model",",","loss","and","optimizer",".","model",",","loss_fn",",","optimizer","=","create_model_loss_and_optimizer","(",")","#","initialize","the","palettizer","config","=","DKMPalettizerConfig","(","global_config=ModuleDKMPalettizerConfig","(","n_bits=4",")",")","palettizer","=","DKMPalettizer","(","model",",","config",")","#","prepare","the","model","to","insert","FakePalettize","layers","for","palettization","model","=","palettizer.prepare","(","inplace=True",")","#","use","palettizer","in","your","PyTorch","training","loop","for","inputs",",","labels","in","data",":","output","=","model","(","inputs",")","loss","=","loss_fn","(","output",",","labels",")","loss.backward","(",")","optimizer.step","(",")","palettizer.step","(",")","#","fold","LUT","and","indices","into","weights","model","=","palettizer.finalize","(","inplace=True",")","Args",":","model","(",":","py",":","class",":","`","torch.nn.Module","`",")",":","Model","on","which","the","palettizer","will","act",".","config","(",":","py",":","class",":","`","DKMPalettizerConfig","`",")",":","Config","which","specifies","how","different","submodules","in","the","model","will","be","configured","for","palettization",".","Default","config","is","used","when","passed","as","``","None","``","."]
docstr: Method to palettize supported modules.
docstr_tok: ["Method","to","palettize","supported","modules","."]
code: def _palettize_supported_modules(self):
      """
      Method to palettize supported modules.
      """
      for name, submodule in self._model.named_modules(
          remove_duplicate=True
      ):
          config = self._config.get_module_config(name, submodule)
          if type(submodule) in self._supported_modules:
              if config is not None:
                  submod_configs = (
                      config if isinstance(config, list) else [config]
                  )
                  for submod_config in submod_configs:
                      if (
                          submodule.weight.numel()
                          > submod_config.weight_threshold
                      ):
                          module_level_advanced_options = (
                              self._get_module_level_advanced_options(
                                  submodule, submod_config
                              )
                          )
                          n_bits = (
                              submod_config.n_bits
                              if submod_config.n_bits is not None
                              else _DEFAULT_PALETTIZATION_SCHEME[
                                  type(submodule)
                              ]["n_bits"]
                          )
                          cluster_dim = (
                              submod_config.cluster_dim
                              if submod_config.cluster_dim is not None
                              else _DEFAULT_PALETTIZATION_SCHEME[
                                  type(submodule)
                              ]["cluster_dim"]
                          )
                          self._palettize_module(
                              submodule,
                              n_bits,
                              cluster_dim,
                              submod_config.quant_min,
                              submod_config.quant_max,
                              submod_config.cluster_dtype,
                              submod_config.dtype,
                              submod_config.quantize_activations,
                              module_level_advanced_options,
                          )
                          self._milestones[
                              name
                          ] = submod_config.milestone
  
imports: import copy
  import logging
  from typing import Dict
  from typing import Optional
  import torch
  import torch.nn
  from torch.ao.quantization import FakeQuantize
  from coremltools.optimize.torch._typing import ParamsDict
  from coremltools.optimize.torch._utils.math_utils import rmse_error
  from coremltools.optimize.torch._utils.torch_utils import get_eval_model
  from coremltools.optimize.torch.base_model_optimizer import BaseModelOptimizer
  from coremltools.optimize.torch.base_model_optimizer import _Report
  from coremltools.optimize.torch.palettization._custom_conversion import PALETTIZATION_CONVERT_DICT
  from coremltools.optimize.torch.palettization._supported_modules import _get_palettization_qat_mappings
  from coremltools.optimize.torch.palettization.fake_palettize import FakePalettize
  from coremltools.optimize.torch.palettization.palettization_config import DEFAULT_PALETTIZATION_ADVANCED_OPTIONS
  from coremltools.optimize.torch.palettization.palettization_config import DEFAULT_PALETTIZATION_SCHEME
  from coremltools.optimize.torch.palettization.palettization_config import DKMPalettizerConfig
  from coremltools.optimize.torch.palettization.palettization_config import ModuleDKMPalettizerConfig
label_desc: Use image node_id 2 for calling the DKMPalettizer obj's underlying member method code with example usage: obj._palettize_supported_modules() without return types