code: def _construct_jagged_tensors_cw(
      embeddings: List[torch.Tensor],
      features: KJTList,
      embedding_names_per_rank: List[List[str]],
      need_indices: bool,
      features_to_permute_indices: Dict[str, torch.Tensor],
      key_to_feature_permuted_coordinates: Dict[str, torch.Tensor],
  ) -> Dict[str, JaggedTensor]:
      ret: Dict[str, JaggedTensor] = {}
      stride = features[0].stride()
      lengths_lists: List[List[torch.Tensor]] = []
      embeddings_lists: List[List[torch.Tensor]] = []
      values_lists: List[List[torch.Tensor]] = []
      for i in range(len(features)):
          embedding = embeddings[i]
          feature = features[i]
          lengths_lists.append(
              torch.unbind(feature.lengths().view(-1, stride), dim=0)
          )
          embeddings_lists.append(
              list(
                  torch.split(
                      embedding, feature.length_per_key(), dim=0
                  )
              )
          )
      if need_indices:
          for i in range(len(features)):
              feature = features[i]
              values_lists.append(
                  list(
                      torch.split(
                          feature.values(), feature.length_per_key()
                      )
                  )
              )
  
      for (
          key,
          permuted_coordinate_tensor,
      ) in key_to_feature_permuted_coordinates.items():
          permuted_coordinates: List[
              List[int]
          ] = permuted_coordinate_tensor.tolist()
  
          rank0, idx_in_rank0 = permuted_coordinates[0]
          ret[key] = JaggedTensor(
              lengths=lengths_lists[rank0][idx_in_rank0],
              values=torch.cat(
                  [
                      embeddings_lists[rank][idx_in_rank]
                      for rank, idx_in_rank in permuted_coordinates
                  ],
                  dim=1,
              ),
              weights=values_lists[rank0][idx_in_rank0]
              if need_indices
              else None,
          )
      return 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 5 for calling a global function with example usage: _construct_jagged_tensors_cw(embeddings, features, embedding_names_per_rank, need_indices, features_to_permute_indices, key_to_feature_permuted_coordinates) and returns: ret