class_docstr: EmbeddingCollection represents a collection of non-pooled embeddings.
  
  It processes sparse data in the form of `KeyedJaggedTensor` of the form [F X B X L]
  where:
  
  * F: features (keys)
  * B: batch size
  * L: length of sparse features (variable)
  
  and outputs `Dict[feature (key), JaggedTensor]`.
  Each `JaggedTensor` contains values of the form (B * L) X D
  where:
  
  * B: batch size
  * L: length of sparse features (jagged)
  * D: each feature's (key's) embedding dimension and lengths are of the form L
  
  Args:
      tables (List[EmbeddingConfig]): list of embedding tables.
      device (Optional[torch.device]): default compute device.
      need_indices (bool): if we need to pass indices to the final lookup result dict
  
  Example::
  
      e1_config = EmbeddingConfig(
          name="t1", embedding_dim=3, num_embeddings=10, feature_names=["f1"]
      )
      e2_config = EmbeddingConfig(
          name="t2", embedding_dim=3, num_embeddings=10, feature_names=["f2"]
      )
  
      ec = EmbeddingCollection(tables=[e1_config, e2_config])
  
      #     0       1        2  <-- batch
      # 0   [0,1] None    [2]
      # 1   [3]    [4]    [5,6,7]
      # ^
      # feature
  
      features = KeyedJaggedTensor.from_offsets_sync(
          keys=["f1", "f2"],
          values=torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]),
          offsets=torch.tensor([0, 2, 2, 3, 4, 5, 8]),
      )
      feature_embeddings = ec(features)
      print(feature_embeddings['f2'].values())
      tensor([[-0.2050,  0.5478,  0.6054],
      [ 0.7352,  0.3210, -3.0399],
      [ 0.1279, -0.1756, -0.4130],
      [ 0.7519, -0.4341, -0.0499],
      [ 0.9329, -1.0697, -0.8095]], grad_fn=<EmbeddingBackward>)
class_docstr_tok: ["EmbeddingCollection","represents","a","collection","of","non-pooled","embeddings",".","It","processes","sparse","data","in","the","form","of","`","KeyedJaggedTensor","`","of","the","form","[","F","X","B","X","L","]","where",":","*","F",":","features","(","keys",")","*","B",":","batch","size","*","L",":","length","of","sparse","features","(","variable",")","and","outputs","`","Dict","[","feature","(","key",")",",","JaggedTensor","]","`",".","Each","`","JaggedTensor","`","contains","values","of","the","form","(","B","*","L",")","X","D","where",":","*","B",":","batch","size","*","L",":","length","of","sparse","features","(","jagged",")","*","D",":","each","feature","'s","(","key","'s",")","embedding","dimension","and","lengths","are","of","the","form","L","Args",":","tables","(","List","[","EmbeddingConfig","]",")",":","list","of","embedding","tables",".","device","(","Optional","[","torch.device","]",")",":","default","compute","device",".","need_indices","(","bool",")",":","if","we","need","to","pass","indices","to","the","final","lookup","result","dict","Example",":",":","e1_config","=","EmbeddingConfig","(","name=","''","t1","''",",","embedding_dim=3",",","num_embeddings=10",",","feature_names=","[","``","f1","''","]",")","e2_config","=","EmbeddingConfig","(","name=","''","t2","''",",","embedding_dim=3",",","num_embeddings=10",",","feature_names=","[","``","f2","''","]",")","ec","=","EmbeddingCollection","(","tables=","[","e1_config",",","e2_config","]",")","#","0","1","2","<","--","batch","#","0","[","0,1","]","None","[","2","]","#","1","[","3","]","[","4","]","[","5,6,7","]","#","^","#","feature","features","=","KeyedJaggedTensor.from_offsets_sync","(","keys=","[","``","f1","''",",","``","f2","''","]",",","values=torch.tensor","(","[","0",",","1",",","2",",","3",",","4",",","5",",","6",",","7","]",")",",","offsets=torch.tensor","(","[","0",",","2",",","2",",","3",",","4",",","5",",","8","]",")",",",")","feature_embeddings","=","ec","(","features",")","print","(","feature_embeddings","[","'f2","'","]",".values","(",")",")","tensor","(","[","[","-0.2050",",","0.5478",",","0.6054","]",",","[","0.7352",",","0.3210",",","-3.0399","]",",","[","0.1279",",","-0.1756",",","-0.4130","]",",","[","0.7519",",","-0.4341",",","-0.0499","]",",","[","0.9329",",","-1.0697",",","-0.8095","]","]",",","grad_fn=","<","EmbeddingBackward",">",")"]
code: def __init__(  # noqa C901
      self,
      tables: List[EmbeddingConfig],
      device: torch.device,
      need_indices: bool = False,
      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._emb_modules: List[
          IntNBitTableBatchedEmbeddingBagsCodegen
      ] = []
      self.embeddings: nn.ModuleDict = nn.ModuleDict()
  
      self._embedding_configs = tables
      self._embedding_dim: int = -1
      self._need_indices: bool = need_indices
      self._output_dtype = output_dtype
      self._device = device
      self.row_alignment = row_alignment
  
      table_names = set()
      for config in tables:
          if config.name in table_names:
              raise ValueError(f"Duplicate table name {config.name}")
          table_names.add(config.name)
          self._embedding_dim = (
              config.embedding_dim
              if self._embedding_dim < 0
              else self._embedding_dim
          )
          if self._embedding_dim != config.embedding_dim:
              raise ValueError(
                  "All tables in a EmbeddingCollection are required to have same embedding dimension."
                  + f" Violating case: {config.name}'s embedding_dim {config.embedding_dim} !="
                  + f" {self._embedding_dim}"
              )
          weight_lists: Optional[
              List[Tuple[torch.Tensor, Optional[torch.Tensor]]]
          ] = ([] if table_name_to_quantized_weights else None)
          if table_name_to_quantized_weights:
              none_throws(weight_lists).append(
                  table_name_to_quantized_weights[config.name]
              )
          emb_module = IntNBitTableBatchedEmbeddingBagsCodegen(
              embedding_specs=[
                  (
                      config.name,
                      config.num_embeddings,
                      config.embedding_dim,
                      data_type_to_sparse_type(config.data_type),
                      EmbeddingLocation.HOST
                      if device.type == "cpu"
                      else EmbeddingLocation.DEVICE,
                  )
              ],
              pooling_mode=PoolingMode.SUM,
              weight_lists=weight_lists,
              device=device,
              output_dtype=data_type_to_sparse_type(
                  dtype_to_data_type(output_dtype)
              ),
              row_alignment=row_alignment,
          )
          if weight_lists is None:
              emb_module.initialize_weights()
  
          self._emb_modules.append(emb_module)
          self.embeddings[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.
          weights_list = (
              emb_module.split_embedding_weights_with_scale_bias(
                  split_scale_bias_mode=2
                  if quant_state_dict_split_scale_bias
                  else 0
              )
          )
  
          weight = weights_list[0][0]
          if config.data_type == DataType.INT4:
              weight = UInt4Tensor(weight)
          elif config.data_type == DataType.INT2:
              weight = UInt2Tensor(weight)
  
          self.embeddings[config.name].register_buffer("weight", weight)
          if quant_state_dict_split_scale_bias:
              self.embeddings[config.name].register_buffer(
                  "weight_qscale", weights_list[0][1]
              )
              self.embeddings[config.name].register_buffer(
                  "weight_qbias", weights_list[0][2]
              )
  
          if not config.feature_names:
              config.feature_names = [config.name]
  
      self._embedding_names_by_table: List[
          List[str]
      ] = get_embedding_names_by_table(tables)
      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 EmbeddingCollection object from inherited base classes: EmbeddingCollectionInterface, ModuleNoCopyMixin with example: obj = EmbeddingCollection(tables, device, need_indices, output_dtype, table_name_to_quantized_weights, register_tbes, quant_state_dict_split_scale_bias, row_alignment)