class_docstr: Sharded implementation of `QuantEmbeddingCollection`.
class_docstr_tok: ["Sharded","implementation","of","`","QuantEmbeddingCollection","`","."]
code: def output_dist(
      self,
      ctx: EmbeddingCollectionContext,
      output: List[List[torch.Tensor]],
  ) -> Dict[str, JaggedTensor]:
      emb_per_sharding: List[List[torch.Tensor]] = []
      features_before_input_dist_per_sharding: List[
          KeyedJaggedTensor
      ] = []
      features_per_sharding: List[KJTList] = []
      unbucketize_tensor_idxs_per_sharding: List[int] = []
      unbucketize_tensors: List[torch.Tensor] = []
      for sharding_output_dist, embeddings, sharding_ctx in zip(
          self._output_dists,
          output,
          ctx.sharding_contexts,
      ):
          sharding_output_dist_res: List[
              torch.Tensor
          ] = sharding_output_dist.forward(embeddings, sharding_ctx)
          emb_per_sharding.append(sharding_output_dist_res)
          features_per_sharding.append(sharding_ctx.features)
          if sharding_ctx.unbucketize_permute_tensor is None:
              unbucketize_tensor_idxs_per_sharding.append(-1)
          else:
              unbucketize_tensors.append(
                  _fx_unwrap_optional_tensor(
                      sharding_ctx.unbucketize_permute_tensor
                  )
              )
              unbucketize_tensor_idxs_per_sharding.append(
                  len(unbucketize_tensors) - 1
              )
          features_before_input_dist_per_sharding.append(
              # pyre-ignore
              sharding_ctx.features_before_input_dist
          )
      return output_jt_dict(
          sharding_types=list(self._sharding_type_to_sharding.keys()),
          emb_per_sharding=emb_per_sharding,
          features_per_sharding=features_per_sharding,
          embedding_names_per_rank_per_sharding=self._embedding_names_per_rank_per_sharding,
          need_indices=self._need_indices,
          features_before_input_dist_per_sharding=features_before_input_dist_per_sharding,
          unbucketize_tensor_idxs_per_sharding=unbucketize_tensor_idxs_per_sharding,
          unbucketize_tensors=unbucketize_tensors,
          features_to_permute_indices=self._features_to_permute_indices,
          key_to_feature_permuted_coordinates_per_sharding=self._key_to_feature_permuted_coordinates_per_sharding,
      )
  
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 13 for calling the ShardedQuantEmbeddingCollection obj's underlying member method code with example usage: obj.output_dist(ctx, output) and returns: output_jt_dict