code: def _construct_jagged_tensors(
      embeddings: torch.Tensor,
      features: KeyedJaggedTensor,
      embedding_names: List[str],
      need_indices: bool = False,
      features_to_permute_indices: Optional[
          Dict[str, List[int]]
      ] = None,
      original_features: Optional[KeyedJaggedTensor] = None,
      reverse_indices: Optional[torch.Tensor] = None,
  ) -> Dict[str, JaggedTensor]:
      with record_function("## _construct_jagged_tensors ##"):
          if original_features is not None:
              features = original_features
          if reverse_indices is not None:
              embeddings = torch.index_select(
                  embeddings, 0, reverse_indices.to(torch.int32)
              )
  
          ret: Dict[str, JaggedTensor] = {}
          stride = features.stride()
          length_per_key = features.length_per_key()
          values = features.values()
  
          lengths = features.lengths().view(-1, stride)
          lengths_tuple = torch.unbind(lengths.view(-1, stride), dim=0)
          embeddings_list = torch.split(
              embeddings, length_per_key, dim=0
          )
          values_list = (
              torch.split(values, length_per_key)
              if need_indices
              else None
          )
  
          key_indices = defaultdict(list)
          for i, key in enumerate(embedding_names):
              key_indices[key].append(i)
          for key, indices in key_indices.items():
              # combines outputs in correct order for CW sharding
              indices = (
                  _permute_indices(
                      indices, features_to_permute_indices[key]
                  )
                  if features_to_permute_indices
                  and key in features_to_permute_indices
                  else indices
              )
              ret[key] = JaggedTensor(
                  lengths=lengths_tuple[indices[0]],
                  values=embeddings_list[indices[0]]
                  if len(indices) == 1
                  else torch.cat(
                      [embeddings_list[i] for i in indices], dim=1
                  ),
                  weights=values_list[indices[0]]
                  if values_list
                  else None,
              )
          return ret
  
imports: import copy
  import logging
  import warnings
  from collections import defaultdict, deque, OrderedDict
  from dataclasses import dataclass, field
  from itertools import accumulate
  from typing import Any, cast, Dict, List, MutableMapping, Optional, Type, Union
  import torch
  from torch import nn
  from torch.autograd.profiler import record_function
  from torch.nn.parallel import DistributedDataParallel
  from torchrec.distributed.embedding_sharding import EmbeddingSharding, EmbeddingShardingInfo, KJTListSplitsAwaitable
  from torchrec.distributed.embedding_types import BaseEmbeddingSharder, EmbeddingComputeKernel, KJTList, ShardedEmbeddingModule, ShardingType
  from torchrec.distributed.sharding.cw_sequence_sharding import CwSequenceEmbeddingSharding
  from torchrec.distributed.sharding.dp_sequence_sharding import DpSequenceEmbeddingSharding
  from torchrec.distributed.sharding.rw_sequence_sharding import RwSequenceEmbeddingSharding
  from torchrec.distributed.sharding.rw_sharding import RwSparseFeaturesDist
  from torchrec.distributed.sharding.sequence_sharding import SequenceShardingContext
  from torchrec.distributed.sharding.tw_sequence_sharding import TwSequenceEmbeddingSharding
  from torchrec.distributed.types import Awaitable, EmbeddingModuleShardingPlan, LazyAwaitable, Multistreamable, ParameterSharding, QuantizedCommCodecs, ShardedTensor, ShardingEnv, ShardMetadata
  from torchrec.distributed.utils import add_params_from_parameter_sharding, convert_to_fbgemm_types, merge_fused_params, optimizer_type_to_emb_opt_type
  from torchrec.modules.embedding_configs import EmbeddingConfig, EmbeddingTableConfig, PoolingType
  from torchrec.modules.embedding_modules import EmbeddingCollection, EmbeddingCollectionInterface
  from torchrec.optim.fused import EmptyFusedOptimizer, FusedOptimizerModule
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer
  from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor
label_desc: Use image node_id 5 for calling a global function with example usage: _construct_jagged_tensors(embeddings, features, embedding_names, need_indices, features_to_permute_indices, original_features, reverse_indices) and returns: ret