code: def _generate_pooled_embedding_batch(
      keys: List[str],
      dims: List[int],
      splits: List[int],
      batch_size_per_rank: List[int],
  ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
      world_size = len(splits)
      offsets = [0] + list(itertools.accumulate(splits))
      local_emb = {}
      B_global = sum(batch_size_per_rank)
      B_offsets = [0] + list(itertools.accumulate(batch_size_per_rank))
  
      for key, dim in zip(keys, dims):
          local_emb[key] = [
              [random.random() for _ in range(dim)]
              for _ in range(B_global)
          ]
  
      in_tensor: List[torch.Tensor] = []
      out_tensor: List[torch.Tensor] = []
      for i in range(world_size):
          in_keys = keys[offsets[i] : offsets[i + 1]]
          in_tensor.append(
              _to_tensor(
                  [
                      local_emb[key][b]
                      for b in range(B_global)
                      for key in in_keys
                  ],
                  torch.float,
              ).view(B_global, -1)
              if in_keys
              else torch.empty(B_global, 0, dtype=torch.float)
          )
          out_tensor.append(
              _to_tensor(
                  [
                      local_emb[key][b]
                      for b in range(B_offsets[i], B_offsets[i + 1])
                      for key in keys
                  ],
                  torch.float,
              ).view(batch_size_per_rank[i], -1)
          )
  
      return in_tensor, 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 5 for calling a global function with example usage: _generate_pooled_embedding_batch(keys, dims, splits, batch_size_per_rank) and returns: in_tensor,  out_tensor