code: def _generate_sequence_embedding_batch(
      keys: List[str],
      dim: int,
      splits: List[int],
      batch_size_per_rank: List[int],
      lengths_before_a2a_per_rank: Dict[int, List],  # pyre-ignore [24]
  ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
      world_size = len(splits)
  
      tensor_by_feature: Dict[
          str, List[torch.Tensor]
      ] = {}  # Model parallel, key as feature
      tensor_by_rank: Dict[
          str, List[torch.Tensor]
      ] = {}  # Data parallel, key as rank
  
      emb_by_rank_feature = {}
      for rank in range(world_size):
          offset = 0
          current_rank_lengths = lengths_before_a2a_per_rank[rank]
          current_rank_batch_size = batch_size_per_rank[rank]
  
          for feature in keys:
              current_stride_lengths = current_rank_lengths[
                  offset : offset + current_rank_batch_size
              ]
              offset += current_rank_batch_size
              emb_by_rank_feature[
                  f"{feature}_{str(rank)}"
              ] = torch.rand(
                  (sum(current_stride_lengths), dim)
              ).tolist()
              tensor_by_feature[f"{feature}"] = []
              tensor_by_rank[f"{str(rank)}"] = []
  
      for k, v in emb_by_rank_feature.items():
          feature, rank = k.split("_")
          tensor_by_feature[feature].extend(v)
          tensor_by_rank[rank].extend(v)
  
      in_tensor: List[torch.Tensor] = []
      out_tensor: List[torch.Tensor] = []
  
      for _, v in tensor_by_feature.items():
          in_tensor.append(torch.Tensor(v))
  
      for _, v in tensor_by_rank.items():
          out_tensor.append(torch.Tensor(v))
  
      input_offsets = [0] + list(itertools.accumulate(splits))
      output_offsets = torch.arange(
          0, world_size + 1, dtype=torch.int
      ).tolist()
  
      regroup_in_tensor: List[torch.Tensor] = []
      regroup_out_tensor: List[torch.Tensor] = []
  
      for i in range(world_size):
          regroup_in_tensor.append(
              torch.cat(
                  in_tensor[input_offsets[i] : input_offsets[i + 1]]
              )
          )
          regroup_out_tensor.append(
              torch.cat(
                  out_tensor[output_offsets[i] : output_offsets[i + 1]]
              )
          )
  
      return regroup_in_tensor, regroup_out_tensor
  
imports: import itertools
  import random
  import unittest
  from typing import Dict, Generator, Iterable, List, Optional, Tuple, TypeVar, Union
  import hypothesis.strategies
  import torch
  import torch.distributed
  from hypothesis import given, settings
  from torchrec.distributed.dist_data import _get_recat, KJTAllToAll, KJTAllToAllSplitsAwaitable, PooledEmbeddingsAllGather, PooledEmbeddingsAllToAll, PooledEmbeddingsReduceScatter, SequenceEmbeddingsAllToAll, VariableBatchPooledEmbeddingsAllToAll
  from torchrec.distributed.fbgemm_qcomm_codec import CommType, get_qcomm_codecs, QCommsConfig
  from torchrec.distributed.test_utils.multi_process import MultiProcessTestBase
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 6 for calling a global function with example usage: _generate_sequence_embedding_batch(keys, dim, splits, batch_size_per_rank, lengths_before_a2a_per_rank) and returns: regroup_in_tensor,  regroup_out_tensor