code: def _construct_jagged_tensors_tw(
      embeddings: List[torch.Tensor],
      features: KJTList,
      need_indices: bool,
  ) -> Dict[str, JaggedTensor]:
      ret: Dict[str, JaggedTensor] = {}
      for i in range(len(embeddings)):
          embeddings_i: torch.Tensor = embeddings[i]
          features_i: KeyedJaggedTensor = features[i]
  
          lengths = features_i.lengths().view(-1, features_i.stride())
          values = features_i.values()
          length_per_key = features_i.length_per_key()
  
          embeddings_list = torch.split(
              embeddings_i, length_per_key, dim=0
          )
          stride = features_i.stride()
          lengths_tuple = torch.unbind(lengths.view(-1, stride), dim=0)
          if need_indices:
              values_list = torch.split(values, length_per_key)
              for i, key in enumerate(features_i.keys()):
                  ret[key] = JaggedTensor(
                      lengths=lengths_tuple[i],
                      values=embeddings_list[i],
                      weights=values_list[i],
                  )
          else:
              for i, key in enumerate(features_i.keys()):
                  ret[key] = JaggedTensor(
                      lengths=lengths_tuple[i],
                      values=embeddings_list[i],
                      weights=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 3 for calling a global function with example usage: _construct_jagged_tensors_tw(embeddings, features, need_indices) and returns: ret