class_docstr: Shards embedding bags column-wise, i.e.. a given embedding table is partitioned
  along its columns and placed on specified ranks.
class_docstr_tok: ["Shards","embedding","bags","column-wise",",","i.e","..","a","given","embedding","table","is","partitioned","along","its","columns","and","placed","on","specified","ranks","."]
code: def create_output_dist(
      self,
      device: Optional[torch.device] = None,
  ) -> BaseEmbeddingDist[
      EmbeddingShardingContext, torch.Tensor, torch.Tensor
  ]:
      device = device if device is not None else self._device
      embedding_permute_op: Optional[
          PermutePooledEmbeddingsSplit
      ] = None
      callbacks: Optional[
          List[Callable[[torch.Tensor], torch.Tensor]]
      ] = None
      if self._permute_embeddings and self._embedding_order != list(
          range(len(self._embedding_order))
      ):
          assert len(self._embedding_order) == len(self._embedding_dims)
          embedding_permute_op = PermutePooledEmbeddingsSplit(
              self._embedding_dims, self._embedding_order, device=device
          )
          callbacks = [embedding_permute_op]
      assert self._pg is not None
      return TwPooledEmbeddingDist(
          pg=self._pg,
          dim_sum_per_rank=self._dim_sum_per_rank(),
          emb_dim_per_rank_per_feature=self._emb_dim_per_rank_per_feature(),
          device=device,
          callbacks=callbacks,
          qcomm_codecs_registry=self.qcomm_codecs_registry,
      )
  
imports: from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar
  import torch
  import torch.distributed
  from fbgemm_gpu.permute_pooled_embedding_modules_split import PermutePooledEmbeddingsSplit
  from torchrec.distributed.dist_data import EmbeddingsAllToOne
  from torchrec.distributed.embedding_lookup import GroupedPooledEmbeddingsLookup, InferGroupedPooledEmbeddingsLookup
  from torchrec.distributed.embedding_sharding import BaseEmbeddingDist, BaseEmbeddingLookup, BaseSparseFeaturesDist, EmbeddingShardingContext, EmbeddingShardingInfo
  from torchrec.distributed.embedding_types import BaseGroupedFeatureProcessor, EmbeddingComputeKernel, KJTList, ShardedEmbeddingTable
  from torchrec.distributed.sharding.tw_sharding import BaseTwEmbeddingSharding, InferTwSparseFeaturesDist, TwPooledEmbeddingDist, TwSparseFeaturesDist
  from torchrec.distributed.types import NullShardingContext, 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 CwPooledEmbeddingSharding obj's underlying member method code with example usage: obj.create_output_dist(device) and returns: TwPooledEmbeddingDist