code: def _run_test_dist(
      cls,
      rank: int,
      world_size: int,
      _input: KeyedJaggedTensor,
      output: KeyedJaggedTensor,
      backend: str,
      splits: List[int],
  ) -> None:
      dist.init_process_group(
          rank=rank, world_size=world_size, backend=backend
      )
      device = torch.device(f"cuda:{rank}")
      if backend == "gloo":
          device = torch.device("cpu")
      _input = _input.to(device=device)
      output = output.to(device=device)
      pg = dist.group.WORLD
      lengths_a2a = KJTAllToAll(
          # pyre-fixme[6]: For 1st param expected `ProcessGroup` but got
          #  `Optional[_distributed_c10d.ProcessGroup]`.
          pg=pg,
          splits=splits,
      )
      cls._validate(lengths_a2a(_input), output)
      dist.destroy_process_group()
  
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 KJTAllToAllTest obj's underlying member method code with example usage: obj._run_test_dist(cls, rank, world_size, _input, output, backend, splits) without return types