class_docstr: Sharded implementation of `EmbeddingBagCollection`.
  This is part of the public API to allow for manual data dist pipelining.
code: def __init__(
      self,
      module: EmbeddingBagCollectionInterface,
      table_name_to_parameter_sharding: Dict[str, ParameterSharding],
      env: ShardingEnv,
      fused_params: Optional[Dict[str, Any]] = None,
      device: Optional[torch.device] = None,
  ) -> None:
      super().__init__()
      self._embedding_bag_configs: List[
          EmbeddingBagConfig
      ] = module.embedding_bag_configs()
      self._sharding_type_to_sharding_infos: Dict[
          str, List[EmbeddingShardingInfo]
      ] = create_sharding_infos_by_sharding(
          module,
          table_name_to_parameter_sharding,
          "embedding_bags.",
          fused_params,
      )
      self._sharding_type_to_sharding: Dict[
          str,
          EmbeddingSharding[
              NullShardingContext,
              KJTList,
              List[torch.Tensor],
              torch.Tensor,
          ],
      ] = {
          sharding_type: create_infer_embedding_bag_sharding(
              sharding_type, embedding_confings, env
          )
          for sharding_type, embedding_confings in self._sharding_type_to_sharding_infos.items()
      }
      self._device = device
      self._is_weighted: bool = module.is_weighted()
      self._input_dists: List[nn.Module] = []
      self._lookups: List[nn.Module] = []
      self._create_lookups(fused_params, device)
  
      # Ensure output dist is set for post processing from an inference runtime (ie. setting device from runtime).
      self._output_dists: torch.nn.ModuleList = torch.nn.ModuleList()
  
      self._embedding_names: List[str] = []
      self._embedding_dims: List[int] = []
      self._feature_splits: List[int] = []
      self._features_order: List[int] = []
  
      # forward pass flow control
      self._has_uninitialized_input_dist: bool = True
      self._has_uninitialized_output_dist: bool = True
      self._has_features_permute: bool = True
  
      tbes: Dict[
          IntNBitTableBatchedEmbeddingBagsCodegen,
          GroupedEmbeddingConfig,
      ] = get_tbes_to_register_from_iterable(self._lookups)
  
      self._tbes_configs: Dict[
          IntNBitTableBatchedEmbeddingBagsCodegen,
          GroupedEmbeddingConfig,
      ] = tbes
  
      # Optional registration of TBEs for model post processing utilities
      if is_fused_param_register_tbe(fused_params):
          self.tbes: torch.nn.ModuleList = torch.nn.ModuleList(
              tbes.keys()
          )
  
      quant_state_dict_split_scale_bias = (
          is_fused_param_quant_state_dict_split_scale_bias(fused_params)
      )
  
      if quant_state_dict_split_scale_bias:
          self._initialize_torch_state(
              tbes=tbes,
              table_name_to_parameter_sharding=table_name_to_parameter_sharding,
              tables_weights_prefix="embedding_bags",
          )
      else:
          table_wise_sharded_only: bool = all(
              [
                  sharding_type == ShardingType.TABLE_WISE.value
                  for sharding_type in self._sharding_type_to_sharding.keys()
              ]
          )
          assert (
              table_wise_sharded_only
          ), "ROW_WISE,COLUMN_WISE shardings can be used only in 'quant_state_dict_split_scale_bias' mode, specify fused_params[FUSED_PARAM_QUANT_STATE_DICT_SPLIT_SCALE_BIAS]=True to __init__ argument"
  
          self.embedding_bags: nn.ModuleDict = nn.ModuleDict()
          for table in self._embedding_bag_configs:
              self.embedding_bags[table.name] = torch.nn.Module()
  
          for _sharding_type, lookup in zip(
              self._sharding_type_to_sharding.keys(), self._lookups
          ):
              lookup_state_dict = lookup.state_dict()
              for key in lookup_state_dict:
                  if key.endswith(".weight"):
                      table_name = key[: -len(".weight")]
                      self.embedding_bags[table_name].register_buffer(
                          "weight", lookup_state_dict[key]
                      )
  
imports: from typing import Any, Dict, List, Optional, Type
  import torch
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import IntNBitTableBatchedEmbeddingBagsCodegen
  from torch import nn
  from torchrec.distributed.embedding_lookup import EmbeddingComputeKernel
  from torchrec.distributed.embedding_sharding import EmbeddingSharding, EmbeddingShardingInfo
  from torchrec.distributed.embedding_types import BaseQuantEmbeddingSharder, FeatureShardingMixIn, GroupedEmbeddingConfig, KJTList, ListOfKJTList
  from torchrec.distributed.embeddingbag import construct_output_kt, create_sharding_infos_by_sharding
  from torchrec.distributed.fused_params import FUSED_PARAM_QUANT_STATE_DICT_SPLIT_SCALE_BIAS, FUSED_PARAM_REGISTER_TBE_BOOL, get_tbes_to_register_from_iterable, is_fused_param_quant_state_dict_split_scale_bias, is_fused_param_register_tbe
  from torchrec.distributed.quant_state import ShardedQuantEmbeddingModuleState
  from torchrec.distributed.sharding.cw_sharding import InferCwPooledEmbeddingSharding
  from torchrec.distributed.sharding.rw_sharding import InferRwPooledEmbeddingSharding
  from torchrec.distributed.sharding.tw_sharding import InferTwEmbeddingSharding
  from torchrec.distributed.types import NullShardedModuleContext, NullShardingContext, ParameterSharding, ShardingEnv, ShardingType
  from torchrec.distributed.utils import copy_to_device
  from torchrec.modules.embedding_configs import data_type_to_sparse_type, dtype_to_data_type, EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollectionInterface
  from torchrec.modules.feature_processor_ import FeatureProcessorsCollection
  from torchrec.quant.embedding_modules import EmbeddingBagCollection, FeatureProcessedEmbeddingBagCollection, MODULE_ATTR_QUANT_STATE_DICT_SPLIT_SCALE_BIAS
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
label_desc: Use image node_id 1 to create a new ShardedQuantEmbeddingBagCollection object from inherited base classes: ShardedQuantEmbeddingModuleState with example: obj = ShardedQuantEmbeddingBagCollection(module, table_name_to_parameter_sharding, env, fused_params, device)