class_docstr: EmbeddingBagCollection represents a collection of pooled embeddings (EmbeddingBags).
  This EmbeddingBagCollection is quantized for lower precision. It relies on fbgemm quantized ops and provides
  table batching.
  
  It processes sparse data in the form of KeyedJaggedTensor
  with values of the form [F X B X L]
  F: features (keys)
  B: batch size
  L: Length of sparse features (jagged)
  
  and outputs a KeyedTensor with values of the form [B * (F * D)]
  where
  F: features (keys)
  D: each feature's (key's) embedding dimension
  B: batch size
  
  Args:
      table_name_to_quantized_weights (Dict[str, Tuple[Tensor, Tensor]]): map of tables to quantized weights
      embedding_configs (List[EmbeddingBagConfig]): list of embedding tables
      is_weighted: (bool): whether input KeyedJaggedTensor is weighted
      device: (Optional[torch.device]): default compute device
  
  Call Args:
      features: KeyedJaggedTensor,
  
  Returns:
      KeyedTensor
  
  Example::
  
      table_0 = EmbeddingBagConfig(
          name="t1", embedding_dim=3, num_embeddings=10, feature_names=["f1"]
      )
      table_1 = EmbeddingBagConfig(
          name="t2", embedding_dim=4, num_embeddings=10, feature_names=["f2"]
      )
      ebc = EmbeddingBagCollection(tables=[eb1_config, eb2_config])
  
      #        0       1        2  <-- batch
      # "f1"   [0,1] None    [2]
      # "f2"   [3]    [4]    [5,6,7]
      #  ^
      # feature
      features = KeyedJaggedTensor(
          keys=["f1", "f2"],
          values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
          offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
      )
  
      ebc.qconfig = torch.quantization.QConfig(
          activation=torch.quantization.PlaceholderObserver.with_args(
              dtype=torch.qint8
          ),
          weight=torch.quantization.PlaceholderObserver.with_args(dtype=torch.qint8),
      )
  
      qebc = QuantEmbeddingBagCollection.from_float(ebc)
      quantized_embeddings = qebc(features)
class_docstr_tok: ["EmbeddingBagCollection","represents","a","collection","of","pooled","embeddings","(","EmbeddingBags",")",".","This","EmbeddingBagCollection","is","quantized","for","lower","precision",".","It","relies","on","fbgemm","quantized","ops","and","provides","table","batching",".","It","processes","sparse","data","in","the","form","of","KeyedJaggedTensor","with","values","of","the","form","[","F","X","B","X","L","]","F",":","features","(","keys",")","B",":","batch","size","L",":","Length","of","sparse","features","(","jagged",")","and","outputs","a","KeyedTensor","with","values","of","the","form","[","B","*","(","F","*","D",")","]","where","F",":","features","(","keys",")","D",":","each","feature","'s","(","key","'s",")","embedding","dimension","B",":","batch","size","Args",":","table_name_to_quantized_weights","(","Dict","[","str",",","Tuple","[","Tensor",",","Tensor","]","]",")",":","map","of","tables","to","quantized","weights","embedding_configs","(","List","[","EmbeddingBagConfig","]",")",":","list","of","embedding","tables","is_weighted",":","(","bool",")",":","whether","input","KeyedJaggedTensor","is","weighted","device",":","(","Optional","[","torch.device","]",")",":","default","compute","device","Call","Args",":","features",":","KeyedJaggedTensor",",","Returns",":","KeyedTensor","Example",":",":","table_0","=","EmbeddingBagConfig","(","name=","''","t1","''",",","embedding_dim=3",",","num_embeddings=10",",","feature_names=","[","``","f1","''","]",")","table_1","=","EmbeddingBagConfig","(","name=","''","t2","''",",","embedding_dim=4",",","num_embeddings=10",",","feature_names=","[","``","f2","''","]",")","ebc","=","EmbeddingBagCollection","(","tables=","[","eb1_config",",","eb2_config","]",")","#","0","1","2","<","--","batch","#","``","f1","''","[","0,1","]","None","[","2","]","#","``","f2","''","[","3","]","[","4","]","[","5,6,7","]","#","^","#","feature","features","=","KeyedJaggedTensor","(","keys=","[","``","f1","''",",","``","f2","''","]",",","values=torch.tensor","(","[","0",",","1",",","2",",","3",",","4",",","5",",","6",",","7","]",")",",","offsets=torch.tensor","(","[","0",",","2",",","2",",","3",",","4",",","5",",","8","]",")",",",")","ebc.qconfig","=","torch.quantization.QConfig","(","activation=torch.quantization.PlaceholderObserver.with_args","(","dtype=torch.qint8",")",",","weight=torch.quantization.PlaceholderObserver.with_args","(","dtype=torch.qint8",")",",",")","qebc","=","QuantEmbeddingBagCollection.from_float","(","ebc",")","quantized_embeddings","=","qebc","(","features",")"]
code: def __init__(
      self,
      tables: List[EmbeddingBagConfig],
      is_weighted: bool,
      device: torch.device,
      output_dtype: torch.dtype = torch.float,
      table_name_to_quantized_weights: Optional[
          Dict[str, Tuple[Tensor, Tensor]]
      ] = None,
      register_tbes: bool = False,
      quant_state_dict_split_scale_bias: bool = False,
      row_alignment: int = DEFAULT_ROW_ALIGNMENT,
  ) -> None:
      super().__init__()
      self._is_weighted = is_weighted
      self._embedding_bag_configs: List[EmbeddingBagConfig] = tables
      self._key_to_tables: Dict[
          Tuple[PoolingType, DataType], List[EmbeddingBagConfig]
      ] = defaultdict(list)
      self._length_per_key: List[int] = []
      # Registering in a List instead of ModuleList because we want don't want them to be auto-registered.
      # Their states will be modified via self.embedding_bags
      self._emb_modules: List[nn.Module] = []
      self._output_dtype = output_dtype
      self._device: torch.device = device
      self._table_name_to_quantized_weights: Optional[
          Dict[str, Tuple[Tensor, Tensor]]
      ] = None
      self.row_alignment = row_alignment
  
      table_names = set()
      for table in self._embedding_bag_configs:
          if table.name in table_names:
              raise ValueError(f"Duplicate table name {table.name}")
          table_names.add(table.name)
          self._length_per_key.extend(
              [table.embedding_dim] * len(table.feature_names)
          )
          key = (table.pooling, table.data_type)
          self._key_to_tables[key].append(table)
  
      self._sum_length_per_key: int = sum(self._length_per_key)
  
      location = (
          EmbeddingLocation.HOST
          if device.type == "cpu"
          else EmbeddingLocation.DEVICE
      )
  
      for key, emb_configs in self._key_to_tables.items():
          (pooling, data_type) = key
          embedding_specs = []
          weight_lists: Optional[
              List[Tuple[torch.Tensor, Optional[torch.Tensor]]]
          ] = ([] if table_name_to_quantized_weights else None)
          feature_table_map: List[int] = []
          index_remappings: List[Optional[torch.Tensor]] = []
          index_remappings_non_none_count: int = 0
  
          for idx, table in enumerate(emb_configs):
              embedding_specs.append(
                  (
                      table.name,
                      table.num_embeddings,
                      table.embedding_dim,
                      data_type_to_sparse_type(data_type),
                      location,
                  )
              )
              if table_name_to_quantized_weights:
                  none_throws(weight_lists).append(
                      table_name_to_quantized_weights[table.name]
                  )
              feature_table_map.extend([idx] * table.num_features())
              index_remappings.append(table.pruning_indices_remapping)
              if table.pruning_indices_remapping is not None:
                  index_remappings_non_none_count += 1
  
          emb_module = IntNBitTableBatchedEmbeddingBagsCodegen(
              embedding_specs=embedding_specs,
              pooling_mode=pooling_type_to_pooling_mode(pooling),
              weight_lists=weight_lists,
              device=device,
              output_dtype=data_type_to_sparse_type(
                  dtype_to_data_type(output_dtype)
              ),
              row_alignment=row_alignment,
              feature_table_map=feature_table_map,
              # pyre-ignore
              index_remapping=index_remappings
              if index_remappings_non_none_count > 0
              else None,
          )
          if weight_lists is None:
              emb_module.initialize_weights()
          self._emb_modules.append(emb_module)
  
      self._embedding_names: List[str] = list(
          itertools.chain(
              *get_embedding_names_by_table(self._embedding_bag_configs)
          )
      )
      # We map over the parameters from FBGEMM backed kernels to the canonical nn.EmbeddingBag
      # representation. This provides consistency between this class and the EmbeddingBagCollection
      # nn.Module API calls (state_dict, named_modules, etc)
      self.embedding_bags: nn.ModuleDict = nn.ModuleDict()
      for (_key, tables), emb_module in zip(
          self._key_to_tables.items(), self._emb_modules
      ):
          for embedding_config, (weight, qscale, qbias) in zip(
              tables,
              emb_module.split_embedding_weights_with_scale_bias(
                  split_scale_bias_mode=2
                  if quant_state_dict_split_scale_bias
                  else 0
              ),
          ):
              self.embedding_bags[
                  embedding_config.name
              ] = torch.nn.Module()
              # register as a buffer so it's exposed in state_dict.
              # TODO: register as param instead of buffer
              # however, since this is only needed for inference, we do not need to expose it as part of parameters.
              # Additionally, we cannot expose uint8 weights as parameters due to autograd restrictions.
  
              if embedding_config.data_type == DataType.INT4:
                  weight = UInt4Tensor(weight)
              elif embedding_config.data_type == DataType.INT2:
                  weight = UInt2Tensor(weight)
  
              self.embedding_bags[
                  embedding_config.name
              ].register_buffer("weight", weight)
              if quant_state_dict_split_scale_bias:
                  self.embedding_bags[
                      embedding_config.name
                  ].register_buffer("weight_qscale", qscale)
                  self.embedding_bags[
                      embedding_config.name
                  ].register_buffer("weight_qbias", qbias)
  
              if embedding_config.pruning_indices_remapping is not None:
                  self.embedding_bags[
                      embedding_config.name
                  ].register_buffer(
                      "index_remappings_array",
                      emb_module.index_remappings_array,
                  )
  
      setattr(
          self,
          MODULE_ATTR_QUANT_STATE_DICT_SPLIT_SCALE_BIAS,
          quant_state_dict_split_scale_bias,
      )
      setattr(self, MODULE_ATTR_REGISTER_TBES_BOOL, register_tbes)
      self.register_tbes = register_tbes
      if register_tbes:
          self.tbes: torch.nn.ModuleList = torch.nn.ModuleList(
              self._emb_modules
          )
  
imports: import copy
  import itertools
  from collections import defaultdict
  from typing import Callable, cast, Dict, List, Optional, Tuple, Type, Union
  import torch
  import torch.nn
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import EmbeddingLocation, IntNBitTableBatchedEmbeddingBagsCodegen, PoolingMode
  from torch import Tensor
  from torchrec.distributed.utils import none_throws
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, DATA_TYPE_NUM_BITS, data_type_to_sparse_type, DataType, dtype_to_data_type, EmbeddingBagConfig, EmbeddingConfig, pooling_type_to_pooling_mode, PoolingType, QuantConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection, EmbeddingBagCollectionInterface, EmbeddingCollection, EmbeddingCollectionInterface, get_embedding_names_by_table
  from torchrec.modules.feature_processor_ import FeatureProcessorsCollection
  from torchrec.modules.fp_embedding_modules import FeatureProcessedEmbeddingBagCollection
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
  from torchrec.tensor_types import UInt2Tensor, UInt4Tensor
  from torchrec.types import ModuleNoCopyMixin
label_desc: Use image node_id 1 to create a new EmbeddingBagCollection object from inherited base classes: EmbeddingBagCollectionInterface, ModuleNoCopyMixin with example: obj = EmbeddingBagCollection(tables, is_weighted, device, output_dtype, table_name_to_quantized_weights, register_tbes, quant_state_dict_split_scale_bias, row_alignment)