code: def test_sequence_embeddings(
      self,
      variable_batch_size: bool,
      qcomms_config: Optional[QCommsConfig],
  ) -> None:
      world_size = 2
      seq_emb_dim = 3
      features = 3
      keys = [f"F{feature}" for feature in range(features)]
  
      if variable_batch_size:
          variable_batch_size = True
          batch_size_per_rank = [3, 2]
  
          feature_num_per_rank = [1, 2]
  
          lengths_before_a2a_per_rank = {
              0: [3, 0, 2, 4, 1, 2, 1, 2, 0],
              1: [4, 3, 1, 0, 5, 0],
          }
  
          lengths_after_a2a_per_rank = [
              torch.tensor(
                  [3, 0, 2, 4, 3], dtype=int
              ),  # pyre-ignore [6]
              torch.tensor(
                  [4, 1, 2, 1, 0, 1, 2, 0, 5, 0],
                  dtype=int,  # pyre-ignore [6]
              ),
          ]
  
          input_splits_per_rank = {}
          output_splits_per_rank = {}
  
          input_splits_per_rank[0] = [
              5,
              10,
          ]  # sum (3,0,2), sum(4, 1, 2, 1, 2, 0)
          input_splits_per_rank[1] = [
              7,
              6,
          ]  # sum (4, 3), sum(1, 0, 5, 0)
          output_splits_per_rank[0] = [5, 7]  # emb input splits
          output_splits_per_rank[1] = [10, 6]  #
      else:
          variable_batch_size = False
          batch_size_per_rank = [2, 2]
  
          feature_num_per_rank = [1, 2]
  
          lengths_before_a2a_per_rank = {
              0: [3, 4, 1, 2, 6, 0],
              1: [4, 0, 2, 3, 1, 2],
          }
          lengths_after_a2a_per_rank = [
              torch.tensor(
                  [[3, 4, 4, 0]], dtype=int
              ),  # pyre-ignore [6]
              torch.tensor(
                  [[1, 2, 2, 3], [6, 0, 1, 2]],
                  dtype=int,  # pyre-ignore [6]
              ),
          ]
  
          input_splits_per_rank = {}
          output_splits_per_rank = {}
  
          input_splits_per_rank[0] = [
              7,
              9,
          ]  # sum (3,4) rank0, sum(2, 6, 5, 0) for rank 1
          input_splits_per_rank[1] = [
              4,
              8,
          ]  # sum (9,0) rank0, sum(7, 8, 1, 5) for rank 1
          output_splits_per_rank[0] = [7, 4]  # emb input splits
          output_splits_per_rank[1] = [9, 8]  #
  
      _input, output = _generate_sequence_embedding_batch(
          keys=keys,
          dim=seq_emb_dim,
          splits=feature_num_per_rank,
          batch_size_per_rank=batch_size_per_rank,
          lengths_before_a2a_per_rank=lengths_before_a2a_per_rank,
      )
  
      kwargs_per_rank = []
      for rank in range(world_size):
          kwargs_per_rank.append(
              {
                  "_input": _input[rank],
                  "output": output[rank],
                  "input_splits": input_splits_per_rank[rank],
                  "output_splits": output_splits_per_rank[rank],
                  "lengths_after_sdd_a2a": lengths_after_a2a_per_rank[
                      rank
                  ],
                  "features_per_rank": feature_num_per_rank,
                  "batch_size_per_rank": batch_size_per_rank,
                  "qcomms_config": qcomms_config,
              }
          )
      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 2 for calling the SeqEmbeddingsAllToAllTest obj's underlying member method code with example usage: obj.test_sequence_embeddings(variable_batch_size, qcomms_config) without return types