code: def _generate_variable_batch_sparse_features_batch(
      keys: List[str],
      splits: List[int],
      batch_size_per_rank_per_feature: List[List[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 i, key in enumerate(keys):
          lengths[key] = [
              [
                  random.randint(0, 10)
                  for _ in range(
                      sum(batch_size_per_rank_per_feature[rank][i])
                  )
              ]
              for rank in range(world_size)
          ]
          values[key] = [
              [
                  random.randint(0, 1000)
                  for _ in range(sum(lengths[key][j]))
              ]
              for j in range(world_size)
          ]
  
          if weights:
              weights[key] = [
                  [random.random() for _ in range(sum(lengths[key][j]))]
                  for j 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,
                  stride_per_key_per_rank=batch_size_per_rank_per_feature[
                      i
                  ],
                  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]]
          key_indices = [keys.index(k) for k in out_keys]
          batch_sizes_by_rank = list(
              zip(*batch_size_per_rank_per_feature)
          )
          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,
                  stride_per_key_per_rank=[
                      list(_flatten(batch_sizes_by_rank[key_idx]))
                      for key_idx in key_indices
                  ],
                  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 4 for calling a global function with example usage: _generate_variable_batch_sparse_features_batch(keys, splits, batch_size_per_rank_per_feature, is_weighted) and returns: in_jagged,  out_jagged