code: def output_jt_dict(
      sharding_types: List[str],
      emb_per_sharding: List[List[torch.Tensor]],
      features_per_sharding: List[KJTList],
      embedding_names_per_rank_per_sharding: List[List[List[str]]],
      need_indices: bool,
      features_before_input_dist_per_sharding: List[KeyedJaggedTensor],
      features_to_permute_indices: Dict[str, torch.Tensor],
      unbucketize_tensors: List[torch.Tensor],
      unbucketize_tensor_idxs_per_sharding: List[int],
      key_to_feature_permuted_coordinates_per_sharding: List[
          Dict[str, torch.Tensor]
      ],
  ) -> Dict[str, JaggedTensor]:
      jt_dict: Dict[str, JaggedTensor] = {}
      for (
          sharding_type,
          emb_sharding,
          features_sharding,
          embedding_names_per_rank,
          unbucketize_tensor_idx,
          features_before_input_dist,
          key_to_feature_permuted_coordinates,
      ) in zip(
          sharding_types,
          emb_per_sharding,
          features_per_sharding,
          embedding_names_per_rank_per_sharding,
          unbucketize_tensor_idxs_per_sharding,
          features_before_input_dist_per_sharding,
          key_to_feature_permuted_coordinates_per_sharding,
      ):
          jt_dict.update(
              _construct_jagged_tensors(
                  sharding_type=sharding_type,
                  embeddings=emb_sharding,
                  features=features_sharding,
                  embedding_names_per_rank=embedding_names_per_rank,
                  features_before_input_dist=features_before_input_dist,
                  need_indices=need_indices,
                  rw_unbucketize_tensor=unbucketize_tensors[
                      unbucketize_tensor_idx
                  ]
                  if unbucketize_tensor_idx != -1
                  else None,
                  cw_features_to_permute_indices=features_to_permute_indices,
                  key_to_feature_permuted_coordinates=key_to_feature_permuted_coordinates,
              )
          )
      return jt_dict
  
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 7 for calling a global function with example usage: output_jt_dict(sharding_types, emb_per_sharding, features_per_sharding, embedding_names_per_rank_per_sharding, need_indices, features_before_input_dist_per_sharding, features_to_permute_indices, unbucketize_tensors, unbucketize_tensor_idxs_per_sharding, key_to_feature_permuted_coordinates_per_sharding) and returns: jt_dict