class_docstr: Redistributes pooled embedding tensor in TWRW fashion by performing a reduce-scatter
  operation row wise on the host level and then an AlltoAll operation table wise on
  the global level.
  
  Args:
      cross_pg (dist.ProcessGroup): global level ProcessGroup for AlltoAll
          communication.
      intra_pg (dist.ProcessGroup): host level ProcessGroup for reduce-scatter
          communication.
      dim_sum_per_node (List[int]): number of features (sum of dimensions) of the
          embedding for each host.
      emb_dim_per_node_per_feature (List[List[int]]):
      device (Optional[torch.device]): device on which buffers will be allocated.
      qcomm_codecs_registry (Optional[Dict[str, QuantizedCommCodecs]]):
class_docstr_tok: ["Redistributes","pooled","embedding","tensor","in","TWRW","fashion","by","performing","a","reduce-scatter","operation","row","wise","on","the","host","level","and","then","an","AlltoAll","operation","table","wise","on","the","global","level",".","Args",":","cross_pg","(","dist.ProcessGroup",")",":","global","level","ProcessGroup","for","AlltoAll","communication",".","intra_pg","(","dist.ProcessGroup",")",":","host","level","ProcessGroup","for","reduce-scatter","communication",".","dim_sum_per_node","(","List","[","int","]",")",":","number","of","features","(","sum","of","dimensions",")","of","the","embedding","for","each","host",".","emb_dim_per_node_per_feature","(","List","[","List","[","int","]","]",")",":","device","(","Optional","[","torch.device","]",")",":","device","on","which","buffers","will","be","allocated",".","qcomm_codecs_registry","(","Optional","[","Dict","[","str",",","QuantizedCommCodecs","]","]",")",":"]
docstr: Performs reduce-scatter pooled operation on pooled embeddings tensor followed by
  AlltoAll pooled operation.
  
  Args:
      local_embs (torch.Tensor): pooled embeddings tensor to distribute.
  
  Returns:
      Awaitable[torch.Tensor]: awaitable of pooled embeddings tensor.
code: def forward(
      self,
      local_embs: torch.Tensor,
      sharding_ctx: Optional[EmbeddingShardingContext] = None,
  ) -> Awaitable[torch.Tensor]:
      """
      Performs reduce-scatter pooled operation on pooled embeddings tensor followed by
      AlltoAll pooled operation.
  
      Args:
          local_embs (torch.Tensor): pooled embeddings tensor to distribute.
  
      Returns:
          Awaitable[torch.Tensor]: awaitable of pooled embeddings tensor.
      """
      if self._intra_dist is None or self._cross_dist is None:
          self._create_output_dist_modules(sharding_ctx)
      local_rank = self._rank % self._intra_pg.size()
      current_node = self._rank // self._intra_pg.size()
      if (
          sharding_ctx is not None
          and sharding_ctx.variable_batch_per_feature
      ):
          (
              batch_size_per_rank_per_feature_by_cross_group,
              batch_size_per_feature_sum_by_cross_group,
          ) = self._preprocess_batch_size_per_rank_per_feature(
              self._intra_pg.size(),
              self._cross_pg.size(),
              sharding_ctx.batch_size_per_rank_per_feature,
          )
          rs_result = cast(
              VariableBatchPooledEmbeddingsReduceScatter,
              self._intra_dist,
          )(
              local_embs,
              batch_size_per_rank_per_feature=batch_size_per_feature_sum_by_cross_group,
              embedding_dims=self._emb_dim_per_node_per_feature[
                  current_node
              ],
          ).wait()
          return cast(
              VariableBatchPooledEmbeddingsAllToAll, self._cross_dist
          )(
              rs_result,
              batch_size_per_rank_per_feature=batch_size_per_rank_per_feature_by_cross_group[
                  local_rank
              ],
              batch_size_per_feature_pre_a2a=sharding_ctx.batch_size_per_feature_pre_a2a,
          )
      elif (
          sharding_ctx is not None
          and len(set(sharding_ctx.batch_size_per_rank)) > 1
      ):
          # preprocess batch_size_per_rank
          (
              batch_size_per_rank_by_cross_group,
              batch_size_sum_by_cross_group,
          ) = self._preprocess_batch_size_per_rank(
              self._intra_pg.size(),
              self._cross_pg.size(),
              sharding_ctx.batch_size_per_rank,
          )
          # Perform ReduceScatterV within one host
          rs_result = cast(
              PooledEmbeddingsReduceScatter, self._intra_dist
          )(
              local_embs, input_splits=batch_size_sum_by_cross_group
          ).wait()
          return cast(PooledEmbeddingsAllToAll, self._cross_dist)(
              rs_result,
              batch_size_per_rank=batch_size_per_rank_by_cross_group[
                  local_rank
              ],
          )
      else:
          return cast(PooledEmbeddingsAllToAll, self._cross_dist)(
              cast(PooledEmbeddingsReduceScatter, self._intra_dist)(
                  local_embs
              ).wait()
          )
  
imports: import itertools
  import math
  from typing import Any, cast, Dict, List, Optional, Tuple, TypeVar, Union
  import torch
  import torch.distributed
  from torchrec.distributed.comm import get_local_size, intra_and_cross_node_pg
  from torchrec.distributed.dist_data import KJTAllToAll, PooledEmbeddingsAllToAll, PooledEmbeddingsReduceScatter, VariableBatchPooledEmbeddingsAllToAll, VariableBatchPooledEmbeddingsReduceScatter
  from torchrec.distributed.embedding_lookup import GroupedPooledEmbeddingsLookup
  from torchrec.distributed.embedding_sharding import BaseEmbeddingDist, BaseEmbeddingLookup, BaseSparseFeaturesDist, bucketize_kjt_before_all2all, EmbeddingSharding, EmbeddingShardingContext, EmbeddingShardingInfo, group_tables
  from torchrec.distributed.embedding_types import BaseGroupedFeatureProcessor, EmbeddingComputeKernel, GroupedEmbeddingConfig, ShardedEmbeddingTable
  from torchrec.distributed.types import Awaitable, CommOp, QuantizedCommCodecs, ShardedTensorMetadata, ShardingEnv, ShardMetadata
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.streamable import Multistreamable
label_desc: Use image node_id 2 for calling the TwRwPooledEmbeddingDist obj's underlying member method code with example usage: obj.forward(local_embs, sharding_ctx) and returns: cast, cast, cast