code: def _run_test_dist(
      cls,
      rank: int,
      world_size: int,
      _input: torch.Tensor,
      output: torch.Tensor,
      backend: str,
      dim_sum_per_rank: List[int],
      batch_size_per_rank: List[int],
      qcomms_config: Optional[QCommsConfig] = None,
  ) -> None:
      dist.init_process_group(
          rank=rank, world_size=world_size, backend=backend
      )
      pg = dist.group.WORLD
      if backend == "gloo":
          device = torch.device("cpu")
      else:
          device = torch.device(f"cuda:{rank}")
      _input = _input.to(device=device)
      output = output.to(device=device)
  
      codecs = get_qcomm_codecs(qcomms_config)
  
      a2a = PooledEmbeddingsAllToAll(
          # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
          #  `Optional[_distributed_c10d.ProcessGroup]`.
          pg=pg,
          dim_sum_per_rank=dim_sum_per_rank,
          device=device,
          codecs=codecs,
      )
      _input.requires_grad = True
      if len(set(batch_size_per_rank)) > 1:
          # variable batch size
          res = a2a(_input, batch_size_per_rank).wait()
      else:
          res = a2a(_input).wait()
      res.backward(res)
  
      atol, rtol = None, None
      if qcomms_config is not None:
          atol, rtol = 0.01, 0.01
          if (
              qcomms_config.forward_precision == CommType.FP8
              or qcomms_config.backward_precision == CommType.FP8
          ):
              atol, rtol = 0.05, 0.05
  
      torch.testing.assert_close(res, output, rtol=rtol, atol=atol)
  
      torch.testing.assert_close(
          _input.cpu().detach().div_(world_size),
          # pyre-ignore
          _input.grad.cpu().detach(),
          atol=atol,
          rtol=rtol,
      )
  
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 1 for calling the PooledEmbeddingsAllToAllTest obj's underlying member method code with example usage: obj._run_test_dist(cls, rank, world_size, _input, output, backend, dim_sum_per_rank, batch_size_per_rank, qcomms_config) without return types