code: def test_pooled_embedding_all_gather(self) -> None:
      world_size = 2
      embeddding_dim = 10
      batch_size = 2
      embeddings = torch.rand((batch_size * world_size, embeddding_dim))
      embeddings_by_rank = list(
          torch.chunk(embeddings, batch_size, dim=0)
      )
      kwargs_per_rank = []
      for rank in range(world_size):
          kwargs_per_rank.append(
              {
                  "input": embeddings_by_rank[rank],
                  "expected_output": embeddings,
              }
          )
  
      self._run_multi_process_test_per_rank(
          callable=self._run_test_dist,
          world_size=world_size,
          kwargs_per_rank=kwargs_per_rank,
      )
  
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 the PooledEmbeddingsAllGatherTest obj's underlying member method code with example usage: obj.test_pooled_embedding_all_gather() without return types