class_docstr: Shards embedding bags table-wise then row-wise.
class_docstr_tok: ["Shards","embedding","bags","table-wise","then","row-wise","."]
code: def create_output_dist(
      self,
      device: Optional[torch.device] = None,
  ) -> BaseEmbeddingDist[
      EmbeddingShardingContext, torch.Tensor, torch.Tensor
  ]:
      return TwRwPooledEmbeddingDist(
          rank=self._rank,
          cross_pg=cast(dist.ProcessGroup, self._cross_pg),
          intra_pg=cast(dist.ProcessGroup, self._intra_pg),
          dim_sum_per_node=self._dim_sum_per_node(),
          emb_dim_per_node_per_feature=self._emb_dim_per_node_per_feature(),
          device=device if device is not None else self._device,
          qcomm_codecs_registry=self.qcomm_codecs_registry,
      )
  
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 3 for calling the TwRwPooledEmbeddingSharding obj's underlying member method code with example usage: obj.create_output_dist(device) and returns: TwRwPooledEmbeddingDist