code: def _test_alltoall_sequence(
      cls,
      rank: int,
      world_size: int,
      backend: str,
  ) -> None:
      dist.init_process_group(
          rank=rank, world_size=world_size, backend=backend
      )
      device = torch.device(f"cuda:{rank}")
      torch.cuda.set_device(device)
  
      ranks = 2
      tables_mp = [[0], [1, 2]]
      lengths_dp = [
          torch.tensor([[1, 2], [1, 1], [2, 1]], dtype=torch.int),
          torch.tensor([[1, 2], [2, 1], [3, 1]], dtype=torch.int),
      ]  # W, T_g, B_l
      lengths_a2a = [
          torch.tensor([[[1, 2]], [[1, 2]]], dtype=torch.int),  # Rank 0
          torch.tensor(
              [
                  [[1, 1], [2, 1]],  # from Rank 0
                  [[2, 1], [3, 1]],  # from rank 1
              ],
              dtype=torch.int,
          ),  # Rank 1
      ]  # W, W, T_l, B_l
      lengths_mp = [
          torch.tensor(
              [
                  [1, 2, 1, 2],
              ],
              dtype=torch.int,
          ),
          torch.tensor([[1, 1, 2, 1], [2, 1, 3, 1]], dtype=torch.int),
      ]  # w, t_l, b_g
      input_seg = list(
          itertools.accumulate([0] + [len(i) for i in tables_mp])
      )
      input_splits = [
          [
              int(
                  lengths_dp[i][
                      input_seg[j] : input_seg[j + 1], :
                  ].sum()
              )
              for j in range(ranks)
          ]
          for i in range(ranks)
      ]
      output_splits = [
          lengths_a2a[i].sum(dim=(1, 2)).tolist() for i in range(ranks)
      ]
      table_dim = 3
      num_features_per_rank = [len(features) for features in tables_mp]
      seq_all2all_forward_recat = []
      for j in range(ranks):
          for i in range(num_features_per_rank[rank]):
              seq_all2all_forward_recat.append(j + i * ranks)
      seq_all2all_forward_recat_tensor = torch.IntTensor(
          seq_all2all_forward_recat
      )
      seq_all2all_backward_recat = []
      for i in range(num_features_per_rank[rank]):
          for j in range(ranks):
              seq_all2all_backward_recat.append(
                  i + j * num_features_per_rank[rank]
              )
  
      seq_all2all_backward_recat_tensor = torch.IntTensor(
          seq_all2all_backward_recat
      )
      input_embeddings = torch.rand(
          int(lengths_mp[rank].sum()),
          table_dim,
          device=device,
          requires_grad=True,
      )
      lengths_after_sparse_data_all2all = torch.IntTensor(
          lengths_mp[rank]
      )
      a2a_req = comm_ops.alltoall_sequence(
          a2a_sequence_embs_tensor=input_embeddings.cuda(),
          forward_recat_tensor=seq_all2all_forward_recat_tensor.cuda(),
          backward_recat_tensor=seq_all2all_backward_recat_tensor.cuda(),
          lengths_after_sparse_data_all2all=lengths_after_sparse_data_all2all.cuda(),
          input_splits=input_splits[rank],
          output_splits=output_splits[rank],
      )
      seq_embs_out = a2a_req.wait()
      seq_embs_out.backward(seq_embs_out)
      grad = input_embeddings.grad
      assert torch.equal(
          input_embeddings.cpu().detach(),
          # pyre-fixme[16]: Optional type has no attribute `cpu`.
          grad.cpu().detach() * world_size,
      )
  
      dist.destroy_process_group()
  
imports: import itertools
  import multiprocessing
  import os
  import unittest
  from typing import Callable
  import torch
  import torch.distributed
  import torchrec.distributed.comm_ops
  from torchrec.distributed.utils import none_throws
  from torchrec.test_utils import get_free_port, seed_and_log
label_desc: Use image node_id 6 for calling the TestAllToAll obj's underlying member method code with example usage: obj._test_alltoall_sequence(cls, rank, world_size, backend) without return types