code: def _quantize_weight(
      state_dict: Dict[str, torch.Tensor],
      data_type: DataType,
  ) -> List[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
      quant_weight_list = []
      for weight in state_dict.values():
          if (
              weight.dtype == torch.float
              or weight.dtype == torch.float16
          ):
              quantized_weights = torch.ops.fbgemm.FloatOrHalfToFusedNBitRowwiseQuantizedSBHalf(
                  weight, DATA_TYPE_NUM_BITS[data_type]
              )
          else:
              raise Exception("Unsupported dtype: {weight.dtype}")
  
          # weight and 4 byte scale shift (2xfp16)
          quant_weight = quantized_weights[:, :-4]
          scale_shift = quantized_weights[:, -4:]
  
          quant_weight_list.append((quant_weight, scale_shift))
      return quant_weight_list
  
imports: import copy
  import logging
  from typing import Any, Dict, Iterator, List, Optional, Tuple
  import torch
  import torch.distributed
  from fbgemm_gpu.split_embedding_configs import SparseType
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import EmbeddingLocation, IntNBitTableBatchedEmbeddingBagsCodegen, PoolingMode, rounded_row_size_in_bytes
  from torchrec.distributed.batched_embedding_kernel import BaseBatchedEmbedding, BaseBatchedEmbeddingBag, BatchedDenseEmbedding, BatchedDenseEmbeddingBag
  from torchrec.distributed.embedding_kernel import BaseEmbedding
  from torchrec.distributed.embedding_types import compute_kernel_to_embedding_location, GroupedEmbeddingConfig
  from torchrec.distributed.fused_params import is_fused_param_quant_state_dict_split_scale_bias, is_fused_param_register_tbe, tbe_fused_params, TBEToRegisterMixIn
  from torchrec.distributed.utils import append_prefix
  from torchrec.modules.embedding_configs import DATA_TYPE_NUM_BITS, data_type_to_sparse_type, DataType, dtype_to_data_type
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling a global function with example usage: _quantize_weight(state_dict, data_type) and returns: quant_weight_list