code: def _run_test_dist(
      cls,
      rank: int,
      world_size: int,
      input: torch.Tensor,
      expected_output: torch.Tensor,
  ) -> None:
      dist.init_process_group(rank=rank, world_size=2, backend="nccl")
      pg = dist.group.WORLD
      input = input.cuda(rank)
      input.requires_grad = True
      # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
      #  `Optional[_distributed_c10d.ProcessGroup]`.
      ag = PooledEmbeddingsAllGather(pg).cuda(rank)
      actual_output = ag(input).wait()
      s = torch.sum(actual_output)
      s.backward()
      cls._validate(actual_output, expected_output, input, world_size)
  
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 2 for calling the PooledEmbeddingsAllGatherTest obj's underlying member method code with example usage: obj._run_test_dist(cls, rank, world_size, input, expected_output) without return types