code: def test_variable_batch_pooled_embeddings(
      self,
      backend: str,
      B: int,
      features: int,
      is_reversed: bool,
      qcomms_config: Optional[QCommsConfig],
  ) -> None:
      world_size = 2
      keys = [f"F{feature}" for feature in range(features)]
      dims = random.sample([8, 16, 32] * features, features)
      rank0_split = random.randint(1, features - 1)
      splits = [rank0_split, features - rank0_split]
      if is_reversed:
          splits.reverse()
      emb_dim_per_rank_per_feature = []
      f_id = 0
      for split in splits:
          emb_dim_per_feature = []
          for _ in range(split):
              emb_dim_per_feature.append(dims[f_id])
              f_id += 1
          emb_dim_per_rank_per_feature.append(emb_dim_per_feature)
  
      batch_size_per_rank_per_feature_pre_a2a = []
      for _ in range(world_size):
          batch_size_per_feature = [
              random.randint(B, B + 4) for _ in keys
          ]
          batch_size_per_rank_per_feature_pre_a2a.append(
              batch_size_per_feature
          )
  
      batch_size_per_rank_per_feature_post_a2a_per_rank = []
      fid = 0
      for i in range(world_size):
          batch_size_per_rank_per_feature_post_a2a = [
              [] for _ in range(world_size)
          ]
          split = splits[i]
          for _ in range(split):
              for j in range(world_size):
                  batch_size_per_rank_per_feature_post_a2a[j].append(
                      batch_size_per_rank_per_feature_pre_a2a[j][fid]
                  )
              fid += 1
          batch_size_per_rank_per_feature_post_a2a_per_rank.append(
              batch_size_per_rank_per_feature_post_a2a
          )
  
      """
          before input dist:
          r_0
          f_0: [1, 2], [3, 4]
          f_1: [5, 6]
          f_2: [1],    [2],   [3]
  
          r_1
          f_0: [1, 2]
          f_1: [5, 6], [3, 4]
          f_2: [1],    [2]
  
          after input dist (splits: [1, 2]):
          r_0
          f_0: [1, 2], [3, 4], [1, 2]
  
          r_1
          f_1: [5, 6], [5, 6], [3, 4]
          f_2: [1], [2], [3], [1], [2]
  
          output layout
          r_0:
          [r_0_f_0_s_0, r_0_f_0_s_1, r_1_f_0_s_0]
  
          r_1:
          [r_0_f_1_s_0, r_0_f_2_s_0, r_0_f_2_s_1, r_0_f_2_s_2,
           r_1_f_1_s_0, r_1_f_1_s_1, r_1_f_2_s_0, r_1_f_2_s_1]
  
          after output dist
          r_0:
          [r_0_f_0_s_0, r_0_f_0_s_1, r_0_f_1_s_0, r_0_f_2_s_0, r_0_f_2_s_1, r_0_f_2_s_2]
  
          r_1:
          [r_1_f_0_s_0, r_1_f_1_s_0, r_1_f_1_s_1, r_1_f_2_s_0, r_1_f_2_s_1]
          """
  
      def _generate_variable_batch_pooled_embedding_batch(
          keys: List[str],
          dims: List[int],
          splits: List[int],
          batch_size_per_rank_per_feature: List[List[int]],
      ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
          world_size = len(splits)
          offsets = [0] + list(itertools.accumulate(splits))
          local_embs = {}
  
          feature_ind = 0
          for key, dim in zip(keys, dims):
              for rank in range(world_size):
                  local_batch_size = batch_size_per_rank_per_feature[
                      rank
                  ][feature_ind]
                  if rank not in local_embs:
                      local_embs[rank] = {}
                  local_embs[rank][key] = torch.rand(
                      dim * local_batch_size, dtype=torch.float
                  )
              feature_ind += 1
  
          in_tensor: List[torch.Tensor] = []
          out_tensor: List[torch.Tensor] = []
          for i in range(world_size):
              in_keys = keys[offsets[i] : offsets[i + 1]]
              input_tensor_list = []
              for rank in range(world_size):
                  input_tensor_list += [
                      local_embs[rank][key] for key in in_keys
                  ]
              input_tensor = torch.cat(input_tensor_list)
              in_tensor.append(input_tensor)
  
              output_tensor = torch.cat(
                  [local_embs[i][key] for key in keys]
              )
              out_tensor.append(output_tensor)
  
          return in_tensor, out_tensor
  
      _input, output = _generate_variable_batch_pooled_embedding_batch(
          keys=keys,
          dims=dims,
          splits=splits,
          batch_size_per_rank_per_feature=batch_size_per_rank_per_feature_pre_a2a,
      )
  
      kwargs_per_rank = []
      for rank in range(world_size):
          kwargs_per_rank.append(
              {
                  "_input": _input[rank],
                  "output": output[rank],
                  "backend": backend,
                  "emb_dim_per_rank_per_feature": emb_dim_per_rank_per_feature,
                  "batch_size_per_rank_per_feature": batch_size_per_rank_per_feature_post_a2a_per_rank[
                      rank
                  ],
                  "batch_size_per_feature_pre_a2a": batch_size_per_rank_per_feature_pre_a2a[
                      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 VariableBatchPooledEmbeddingsAllToAllTest obj's underlying member method code with example usage: obj.test_variable_batch_pooled_embeddings(backend, B, features, is_reversed, qcomms_config) and returns: in_tensor,  out_tensor