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","]","]",")",":"]
code: def __init__(
      self,
      rank: int,
      cross_pg: dist.ProcessGroup,
      intra_pg: dist.ProcessGroup,
      dim_sum_per_node: List[int],
      emb_dim_per_node_per_feature: List[List[int]],
      device: Optional[torch.device] = None,
      qcomm_codecs_registry: Optional[
          Dict[str, QuantizedCommCodecs]
      ] = None,
  ) -> None:
      super().__init__()
      self._rank = rank
      self._intra_pg: dist.ProcessGroup = intra_pg
      self._cross_pg: dist.ProcessGroup = cross_pg
      self._dim_sum_per_node = dim_sum_per_node
      self._emb_dim_per_node_per_feature = emb_dim_per_node_per_feature
      self._device = device
      self._intra_codecs: Optional[QuantizedCommCodecs] = (
          qcomm_codecs_registry.get(
              CommOp.POOLED_EMBEDDINGS_REDUCE_SCATTER.name, None
          )
          if qcomm_codecs_registry
          else None
      )
      self._cross_codecs: Optional[QuantizedCommCodecs] = (
          qcomm_codecs_registry.get(
              CommOp.POOLED_EMBEDDINGS_ALL_TO_ALL.name, None
          )
          if qcomm_codecs_registry
          else None
      )
      self._intra_dist: Optional[
          Union[
              PooledEmbeddingsReduceScatter,
              VariableBatchPooledEmbeddingsReduceScatter,
          ]
      ] = None
      self._cross_dist: Optional[
          Union[
              PooledEmbeddingsAllToAll,
              VariableBatchPooledEmbeddingsAllToAll,
          ]
      ] = None
  
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 1 to create a new TwRwPooledEmbeddingDist object from inherited base classes: BaseEmbeddingDist with example: obj = TwRwPooledEmbeddingDist(rank, cross_pg, intra_pg, dim_sum_per_node, emb_dim_per_node_per_feature, device, qcomm_codecs_registry)