class_docstr: Sharded implementation of `QuantEmbeddingCollection`.
class_docstr_tok: ["Sharded","implementation","of","`","QuantEmbeddingCollection","`","."]
code: def input_dist(
      self,
      ctx: EmbeddingCollectionContext,
      features: KeyedJaggedTensor,
  ) -> ListOfKJTList:
      if self._has_uninitialized_input_dist:
          self._create_input_dist(
              input_feature_names=features.keys()
              if features is not None
              else [],
              device=features.device(),
              input_dist_device=self._device,
          )
          self._has_uninitialized_input_dist = False
      if self._has_uninitialized_output_dist:
          self._create_output_dist(features.device())
          self._has_uninitialized_output_dist = False
      ret: List[KJTList] = []
      with torch.no_grad():
          features_by_sharding = []
          if self._features_order:
              features = features.permute(
                  self._features_order,
                  self._features_order_tensor,
              )
          features_by_sharding = (
              [features]
              if len(self._feature_splits) == 1
              else features.split(self._feature_splits)
          )
  
          for i in range(len(self._input_dists)):
              input_dist = self._input_dists[i]
              input_dist_result = input_dist.forward(
                  features_by_sharding[i]
              )
              ret.append(input_dist_result)
              ctx.sharding_contexts.append(
                  InferSequenceShardingContext(
                      features=input_dist_result,
                      features_before_input_dist=features,
                      unbucketize_permute_tensor=input_dist.unbucketize_permute_tensor
                      if isinstance(
                          input_dist, InferRwSparseFeaturesDist
                      )
                      else None,
                  )
              )
      return ListOfKJTList(ret)
  
imports: from collections import defaultdict, deque
  from dataclasses import dataclass
  from typing import Any, cast, Dict, List, Optional, Tuple, Type
  import torch
  from fbgemm_gpu.split_table_batched_embeddings_ops_inference import IntNBitTableBatchedEmbeddingBagsCodegen
  from torch import nn
  from torchrec.distributed.embedding import create_sharding_infos_by_sharding, EmbeddingShardingInfo
  from torchrec.distributed.embedding_sharding import EmbeddingSharding
  from torchrec.distributed.embedding_types import BaseQuantEmbeddingSharder, FeatureShardingMixIn, GroupedEmbeddingConfig, KJTList, ListOfKJTList, ShardingType
  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_sequence_sharding import InferCwSequenceEmbeddingSharding
  from torchrec.distributed.sharding.rw_sequence_sharding import InferRwSequenceEmbeddingSharding
  from torchrec.distributed.sharding.rw_sharding import InferRwSparseFeaturesDist
  from torchrec.distributed.sharding.sequence_sharding import InferSequenceShardingContext
  from torchrec.distributed.sharding.tw_sequence_sharding import InferTwSequenceEmbeddingSharding
  from torchrec.distributed.types import ParameterSharding, ShardingEnv, ShardMetadata
  from torchrec.modules.embedding_configs import data_type_to_sparse_type, dtype_to_data_type, EmbeddingConfig
  from torchrec.quant.embedding_modules import EmbeddingCollection, MODULE_ATTR_QUANT_STATE_DICT_SPLIT_SCALE_BIAS
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
  from torchrec.streamable import Multistreamable
label_desc: Use image node_id 10 for calling the ShardedQuantEmbeddingCollection obj's underlying member method code with example usage: obj.input_dist(ctx, features) and returns: ListOfKJTList