class_docstr: EmbeddingCollection represents a unsharded collection of non-pooled embeddings. The semantics
  of this module is that during the backwards pass, the registered optimizer will be called.
  
  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 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__(
      self,
      tables: List[EmbeddingConfig],
      optimizer_type: Type[torch.optim.Optimizer],
      optimizer_kwargs: Dict[str, Any],
      device: Optional[torch.device] = None,
      need_indices: bool = False,
      location: Optional[EmbeddingLocation] = None,
  ) -> None:
      super().__init__()
  
      self._optimizer_type = optimizer_type
      self._optimizer_kwargs = optimizer_kwargs
  
      emb_optim_and_kwargs = convert_optimizer_type_and_kwargs(
          optimizer_type, optimizer_kwargs
      )
      if emb_optim_and_kwargs is None:
          raise ValueError(
              f"Cannot fuse optimizer_type={optimizer_type} with kwargs {optimizer_kwargs}"
          )
      (emb_optim_type, emb_opt_kwargs) = emb_optim_and_kwargs
  
      if location in [
          EmbeddingLocation.DEVICE,
          EmbeddingLocation.MANAGED,
          EmbeddingLocation.MANAGED_CACHING,
      ]:
          assert device is not None and device.type in [
              "cuda",
              "meta",
          ], f"Using location={location} requires device=cuda or meta"
  
      if device is None:
          device = torch.device("cpu")
  
      assert device.type in [
          "cuda",
          "meta",
      ], "FusedEmbeddingCollection is only supported for device in [CUDA, meta] currently. There are plans to support device=CPU."
  
      if location is None:
          if device.type in ["cpu", "meta"]:
              location = EmbeddingLocation.HOST
          elif device.type == "cuda":
              location = EmbeddingLocation.DEVICE
          else:
              raise ValueError("EmbeddingLocation could not be set")
  
      self._embedding_configs = tables
      self._need_indices: bool = need_indices
      self._embedding_dim: int = -1
  
      # 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._key_to_tables: Dict[
          DataType, List[EmbeddingConfig]
      ] = defaultdict(list)
  
      seen_features = set()
      self._shared_features: Set[str] = set()
      for table in tables:
          key = table.data_type
          self._key_to_tables[key].append(table)
  
          if self._embedding_dim == -1:
              self._embedding_dim = table.embedding_dim
          elif self._embedding_dim != table.embedding_dim:
              raise ValueError(
                  "All tables in a EmbeddingCollection are required to have same embedding dimension."
                  + f" Violating case: {table}'s embedding_dim {table.embedding_dim} !="
                  + f" {self._embedding_dim}"
              )
          for feature in table.feature_names:
              if feature in seen_features:
                  self._shared_features.add(feature)
              else:
                  seen_features.add(feature)
  
      optims = []
      for key, tables in self._key_to_tables.items():
          data_type = key
          emb_module = _BatchedFusedEmbeddingLookups(
              cast(List[BaseEmbeddingConfig], tables),
              data_type=data_type,
              pooling=PoolingType.NONE,
              optimizer_type=emb_optim_type,
              optimizer_kwargs=emb_opt_kwargs,
              device=device,
              embedding_location=location,
          )
          self._emb_modules.append(emb_module)
          params: Dict[str, torch.Tensor] = {}
          for (
              param_key,
              weight,
          ) in emb_module.fused_optimizer().params.items():
              params[f"embeddings.{param_key}"] = weight
          optims.append(("", emb_module.fused_optimizer()))
  
      self._optim: CombinedOptimizer = CombinedOptimizer(optims)
      self._embedding_names: List[str] = list(
          itertools.chain(
              *get_embedding_names_by_table(self._embedding_configs)
          )
      )
  
      self._embedding_names_by_table: List[
          List[str]
      ] = get_embedding_names_by_table(
          self._embedding_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's
      # nn.Module API calls (state_dict, named_modules, etc)
      self.embeddings: nn.ModuleDict = nn.ModuleDict()
      for (_key, tables), emb_module in zip(
          self._key_to_tables.items(), self._emb_modules
      ):
          for embedding_config, weight in zip(
              tables,
              emb_module.split_embedding_weights(),
              #  torch._tensor.Tensor]` is not a function.
          ):
              self.embeddings[embedding_config.name] = torch.nn.Module()
              self.embeddings[embedding_config.name].register_parameter(
                  "weight", torch.nn.Parameter(weight)
              )
  
imports: import copy
  import itertools
  from collections import defaultdict
  from typing import Any, cast, Dict, Iterator, List, Optional, Set, Tuple, Type
  import torch
  import torch.nn
  import torchrec.optim
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from fbgemm_gpu.split_table_batched_embeddings_ops_training import ComputeDevice, EmbeddingLocation, SplitTableBatchedEmbeddingBagsCodegen
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, DataType, EmbeddingBagConfig, EmbeddingConfig, pooling_type_to_pooling_mode, PoolingType
  from torchrec.modules.embedding_modules import EmbeddingBagCollection, EmbeddingBagCollectionInterface, EmbeddingCollection, EmbeddingCollectionInterface, get_embedding_names_by_table
  from torchrec.optim.fused import FusedOptimizer, FusedOptimizerModule
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor, KeyedTensor
label_desc: Use image node_id 1 to create a new FusedEmbeddingCollection object from inherited base classes: EmbeddingCollectionInterface, FusedOptimizerModule with example: obj = FusedEmbeddingCollection(tables, optimizer_type, optimizer_kwargs, device, need_indices, location)