code: def _construct_jagged_tensors(
      sharding_type: str,
      embeddings: List[torch.Tensor],
      features: KJTList,
      embedding_names_per_rank: List[List[str]],
      features_before_input_dist: KeyedJaggedTensor,
      need_indices: bool,
      rw_unbucketize_tensor: Optional[torch.Tensor],
      cw_features_to_permute_indices: Dict[str, torch.Tensor],
      key_to_feature_permuted_coordinates: Dict[str, torch.Tensor],
  ) -> Dict[str, JaggedTensor]:
      # Validating sharding type and parameters
      valid_sharding_types = [
          ShardingType.ROW_WISE.value,
          ShardingType.COLUMN_WISE.value,
          ShardingType.TABLE_WISE.value,
      ]
      if sharding_type not in valid_sharding_types:
          raise ValueError(f"Unknown sharding type {sharding_type}")
  
      if (
          sharding_type == ShardingType.ROW_WISE.value
          and rw_unbucketize_tensor is None
      ):
          raise ValueError(
              "rw_unbucketize_tensor is required for row-wise sharding"
          )
  
      if (
          sharding_type == ShardingType.ROW_WISE.value
          and rw_unbucketize_tensor is not None
      ):
          return _construct_jagged_tensors_rw(
              embeddings,
              features_before_input_dist,
              need_indices,
              rw_unbucketize_tensor,
          )
      elif sharding_type == ShardingType.COLUMN_WISE.value:
          return _construct_jagged_tensors_cw(
              embeddings,
              features,
              embedding_names_per_rank,
              need_indices,
              cw_features_to_permute_indices,
              key_to_feature_permuted_coordinates,
          )
      else:  # sharding_type == ShardingType.TABLE_WISE.value
          return _construct_jagged_tensors_tw(
              embeddings, features, need_indices
          )
  
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 6 for calling a global function with example usage: _construct_jagged_tensors(sharding_type, embeddings, features, embedding_names_per_rank, features_before_input_dist, need_indices, rw_unbucketize_tensor, cw_features_to_permute_indices, key_to_feature_permuted_coordinates) and returns: _construct_jagged_tensors_rw, _construct_jagged_tensors_cw, _construct_jagged_tensors_tw