code: def _generate_sparse_features_batch(
      keys: List[str],
      splits: List[int],
      batch_size_per_rank: List[int],
      is_weighted: bool = False,
  ) -> Tuple[List[KeyedJaggedTensor], List[KeyedJaggedTensor]]:
      world_size = len(splits)
      offsets = [0] + list(itertools.accumulate(splits))
      values = {}
      lengths = {}
      weights = {} if is_weighted else None
  
      for key in keys:
          lengths[key] = [
              [
                  random.randint(0, 10)
                  for _ in range(batch_size_per_rank[i])
              ]
              for i in range(world_size)
          ]
          values[key] = [
              [
                  random.randint(0, 1000)
                  for _ in range(sum(lengths[key][i]))
              ]
              for i in range(world_size)
          ]
  
          if weights:
              weights[key] = [
                  [random.random() for _ in range(sum(lengths[key][i]))]
                  for i in range(world_size)
              ]
  
      in_jagged: List[KeyedJaggedTensor] = []
      out_jagged: List[KeyedJaggedTensor] = []
      for i in range(world_size):
          in_jagged.append(
              KeyedJaggedTensor.from_lengths_sync(
                  keys=keys,
                  lengths=_to_tensor(
                      [lengths[key][i] for key in keys], torch.int
                  ),
                  values=_to_tensor(
                      [values[key][i] for key in keys], torch.int
                  ),
                  weights=_to_tensor(
                      [weights[key][i] for key in keys], torch.float
                  )
                  if weights
                  else None,
              )
          )
          key_index = []
          out_keys = keys[offsets[i] : offsets[i + 1]]
          for key in out_keys:
              for j in range(world_size):
                  key_index.append((key, j))
          out_jagged.append(
              KeyedJaggedTensor.from_lengths_sync(
                  keys=out_keys,
                  lengths=_to_tensor(
                      [lengths[key][j] for key, j in key_index],
                      torch.int,
                  ),
                  values=_to_tensor(
                      [values[key][j] for key, j in key_index],
                      torch.int,
                  ),
                  weights=_to_tensor(
                      [weights[key][j] for key, j in key_index],
                      torch.float,
                  )
                  if weights
                  else None,
              )
          )
      return in_jagged, out_jagged
  
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 3 for calling a global function with example usage: _generate_sparse_features_batch(keys, splits, batch_size_per_rank, is_weighted) and returns: in_jagged,  out_jagged