class_docstr: FusedEmbeddingBagCollection represents a collection of pooled embeddings (`EmbeddingBags`).
  It utilizes a technique called Optimizer fusion (register the optimizer with model). The semantics
  of this is that during the backwards pass, the registered optimizer will be called.
  
  It processes sparse data in the form of `KeyedJaggedTensor` with values of the form
  [F X B X L] where:
  
  * F: features (keys)
  * B: batch size
  * L: length of sparse features (jagged)
  
  and outputs a `KeyedTensor` with values of the form [B x F x D] where:
  
  * F: features (keys)
  * D: each feature's (key's) embedding dimension
  * B: batch size
  
  Args:
      tables (List[EmbeddingBagConfig]): list of embedding tables.
      is_weighted (bool): whether input `KeyedJaggedTensor` is weighted.
      optimizer (Type[torch.optim.Optimizer]): fusion optimizer type
      optimizer_kwargs: Dict[str, Any]: fusion optimizer kwargs
      device (Optional[torch.device]): compute device.
  
  Example::
  
      table_0 = EmbeddingBagConfig(
          name="t1", embedding_dim=4, num_embeddings=10, feature_names=["f1"]
      )
      table_1 = EmbeddingBagConfig(
          name="t2", embedding_dim=8, num_embeddings=10, feature_names=["f2"]
      )
  
      ebc = FusedEmbeddingBagCollection(tables=[table_0, table_1], optimizer_type=torch.optim.SGD, optimizer_kwargs={"lr": .01})
  
      #        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]),
      )
  
      pooled_embeddings = ebc(features)
      print(pooled_embeddings.values())
      tensor([[ 0.2093,  0.1395,  0.1571,  0.3583,  0.0421,  0.0037, -0.0692,  0.0663,
        0.2166, -0.3150, -0.2771, -0.0301],
      [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0165, -0.1225,  0.2483,  0.0624,
       -0.1168, -0.0509, -0.1309,  0.3059],
      [ 0.0811, -0.1779, -0.1443,  0.1097, -0.4410, -0.4036,  0.4458, -0.2735,
       -0.3080, -0.2102, -0.0564,  0.5583]], grad_fn=<CatBackward0>)
      print(pooled_embeddings.keys())
      ['f1', 'f2']
      print(pooled_embeddings.offset_per_key())
      [0, 4, 12]
code: def __init__(
      self,
      tables: List[EmbeddingBagConfig],
      optimizer_type: Type[torch.optim.Optimizer],
      optimizer_kwargs: Dict[str, Any],
      is_weighted: bool = False,
      device: Optional[torch.device] = None,
      location: Optional[EmbeddingLocation] = None,
  ) -> None:
      super().__init__()
  
      self._optimizer_type = optimizer_type
      self._optimizer_kwargs = optimizer_kwargs
      self._device: torch.device = (
          device if device is not None else torch.device("cpu")
      )
  
      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")
  
      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._is_weighted = is_weighted
      self._embedding_bag_configs = tables
  
      # 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[
          Tuple[PoolingType, DataType], List[EmbeddingBagConfig]
      ] = defaultdict(list)
  
      self._length_per_key: List[int] = []
  
      for table in tables:
          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)
  
      optims = []
      for key, tables in self._key_to_tables.items():
          (pooling, data_type) = key
          emb_module = _BatchedFusedEmbeddingLookups(
              cast(List[BaseEmbeddingConfig], tables),
              data_type=data_type,
              pooling=pooling,
              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"embedding_bags.{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_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's
      # 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 in zip(
              tables,
              emb_module.split_embedding_weights(),
              #  torch._tensor.Tensor]` is not a function.
          ):
              self.embedding_bags[
                  embedding_config.name
              ] = torch.nn.Module()
              self.embedding_bags[
                  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 FusedEmbeddingBagCollection object from inherited base classes: EmbeddingBagCollectionInterface, FusedOptimizerModule with example: obj = FusedEmbeddingBagCollection(tables, optimizer_type, optimizer_kwargs, is_weighted, device, location)