class_docstr: Redistributes pooled embedding tensor in RW fashion by performing a reduce-scatter
  operation.
  
  Args:
      pg (dist.ProcessGroup): ProcessGroup for reduce-scatter communication.
docstr: Performs reduce-scatter pooled operation on pooled embeddings tensor.
  
  Args:
      local_embs (torch.Tensor): pooled embeddings tensor to distribute.
      sharding_ctx (Optional[EmbeddingShardingContext]): shared context from
          KJTAllToAll operation.
  
  Returns:
      Awaitable[torch.Tensor]: awaitable of pooled embeddings tensor.
docstr_tok: ["Performs","reduce-scatter","pooled","operation","on","pooled","embeddings","tensor",".","Args",":","local_embs","(","torch.Tensor",")",":","pooled","embeddings","tensor","to","distribute",".","sharding_ctx","(","Optional","[","EmbeddingShardingContext","]",")",":","shared","context","from","KJTAllToAll","operation",".","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.
  
      Args:
          local_embs (torch.Tensor): pooled embeddings tensor to distribute.
          sharding_ctx (Optional[EmbeddingShardingContext]): shared context from
              KJTAllToAll operation.
  
      Returns:
          Awaitable[torch.Tensor]: awaitable of pooled embeddings tensor.
      """
      if self._dist is None:
          self._create_output_dist_module(sharding_ctx)
  
      if sharding_ctx is None:
          return cast(PooledEmbeddingsReduceScatter, self._dist)(
              local_embs
          )
      elif sharding_ctx.variable_batch_per_feature:
          return cast(
              VariableBatchPooledEmbeddingsReduceScatter, self._dist
          )(
              local_embs,
              batch_size_per_rank_per_feature=sharding_ctx.batch_size_per_rank_per_feature,
              embedding_dims=self._embedding_dims,
          )
      else:
          return cast(PooledEmbeddingsReduceScatter, self._dist)(
              local_embs,
              input_splits=sharding_ctx.batch_size_per_rank,
          )
  
imports: import math
  from typing import Any, cast, Dict, List, Optional, Tuple, TypeVar, Union
  import torch
  import torch.distributed
  from torchrec.distributed.dist_data import EmbeddingsAllToOneReduce, KJTAllToAll, KJTOneToAll, PooledEmbeddingsReduceScatter, VariableBatchPooledEmbeddingsReduceScatter
  from torchrec.distributed.embedding_lookup import GroupedPooledEmbeddingsLookup, InferGroupedPooledEmbeddingsLookup
  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, KJTList, ShardedEmbeddingTable
  from torchrec.distributed.types import Awaitable, CommOp, NullShardingContext, 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 RwPooledEmbeddingDist obj's underlying member method code with example usage: obj.forward(local_embs, sharding_ctx) and returns: cast, cast, cast