code: def _get_palettization_qat_mappings():
      """
      _get_palettization_qat_mappings creates qat_module_mappings supported by coremltools.optimize.torch for palettization. We
      support three modules already in DEFAULT_QAT_MODULE_MAPPINGS, namely, nn.Linear, nn.Conv2d and
      nn.Conv3d. Additionally, we have added support for preparation of nn.Conv1d, nn.LayerNorm,
      nn.MultiheadAttention and nn.Embedding modules.
      """
      qat_module_mappings = (
          _torch.quantization.quantization_mappings.get_default_qat_module_mappings()
      )
      for k in list(qat_module_mappings.keys()):
          if k not in SUPPORTED_PYTORCH_QAT_MODULES:
              del qat_module_mappings[k]
      qat_module_mappings[Conv1d._FLOAT_MODULE] = Conv1d
      qat_module_mappings[LayerNorm._FLOAT_MODULE] = LayerNorm
      qat_module_mappings[
          MultiheadAttention._FLOAT_MODULE
      ] = MultiheadAttention
      qat_module_mappings[Embedding._FLOAT_MODULE] = Embedding
  
      return qat_module_mappings
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from .palettization_config import SUPPORTED_PYTORCH_QAT_MODULES
label_desc: Use image node_id 1 for calling a global function with example usage: _get_palettization_qat_mappings() and returns: qat_module_mappings