code: def __init__(
      self,
      config: GroupedEmbeddingConfig,
      pg: Optional[dist.ProcessGroup] = None,
      device: Optional[torch.device] = None,
      fused_params: Optional[Dict[str, Any]] = None,
  ) -> None:
      super().__init__(config, pg, device)
  
      managed: List[EmbeddingLocation] = []
      for table in config.embedding_tables:
          if device is not None and device.type == "cuda":
              managed.append(
                  compute_kernel_to_embedding_location(
                      table.compute_kernel
                  )
              )
          else:
              managed.append(EmbeddingLocation.HOST)
      self._config: GroupedEmbeddingConfig = config
      self._emb_module_registered: bool = is_fused_param_register_tbe(
          fused_params
      )
      self._quant_state_dict_split_scale_bias: bool = (
          is_fused_param_quant_state_dict_split_scale_bias(fused_params)
      )
      self._emb_module: IntNBitTableBatchedEmbeddingBagsCodegen = IntNBitTableBatchedEmbeddingBagsCodegen(
          embedding_specs=[
              (
                  table.name,
                  local_rows,
                  local_cols
                  if self._quant_state_dict_split_scale_bias
                  else table.embedding_dim,
                  data_type_to_sparse_type(config.data_type),
                  location,
              )
              for local_rows, local_cols, table, location in zip(
                  self._local_rows,
                  self._local_cols,
                  config.embedding_tables,
                  managed,
              )
          ],
          device=device,
          pooling_mode=PoolingMode.NONE,
          feature_table_map=self._feature_table_map,
          row_alignment=16,
          uvm_host_mapped=True,  # Use cudaHostAlloc for UVM CACHING to fix imbalance numa memory issue
          **(tbe_fused_params(fused_params) or {}),
      )
      if device is not None:
          self._emb_module.initialize_weights()
  
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 1 to create a new QuantBatchedEmbedding object from inherited base classes: BaseBatchedEmbedding, TBEToRegisterMixIn with example: obj = QuantBatchedEmbedding(config, pg, device, fused_params)