class_docstr: Redistributes pooled embedding tensor with an AlltoAll collective operation for
  table wise sharding.
  
  Args:
      pg (dist.ProcessGroup): ProcessGroup for AlltoAll communication.
      dim_sum_per_rank (List[int]): number of features (sum of dimensions) of the
          embedding in each rank.
      emb_dim_per_rank_per_feature (List[List[int]]): embedding dimension per rank per
          feature, used for variable batch per feature.
      device (Optional[torch.device]): device on which buffers will be allocated.
      callbacks (Optional[List[Callable[[torch.Tensor], torch.Tensor]]]):
      qcomm_codecs_registry (Optional[Dict[str, QuantizedCommCodecs]]):
class_docstr_tok: ["Redistributes","pooled","embedding","tensor","with","an","AlltoAll","collective","operation","for","table","wise","sharding",".","Args",":","pg","(","dist.ProcessGroup",")",":","ProcessGroup","for","AlltoAll","communication",".","dim_sum_per_rank","(","List","[","int","]",")",":","number","of","features","(","sum","of","dimensions",")","of","the","embedding","in","each","rank",".","emb_dim_per_rank_per_feature","(","List","[","List","[","int","]","]",")",":","embedding","dimension","per","rank","per","feature",",","used","for","variable","batch","per","feature",".","device","(","Optional","[","torch.device","]",")",":","device","on","which","buffers","will","be","allocated",".","callbacks","(","Optional","[","List","[","Callable","[","[","torch.Tensor","]",",","torch.Tensor","]","]","]",")",":","qcomm_codecs_registry","(","Optional","[","Dict","[","str",",","QuantizedCommCodecs","]","]",")",":"]
code: def _create_output_dist_module(
      self, sharding_ctx: Optional[EmbeddingShardingContext] = None
  ) -> None:
      if (
          sharding_ctx is not None
          and sharding_ctx.variable_batch_per_feature
      ):
          self._dist = VariableBatchPooledEmbeddingsAllToAll(
              pg=self._pg,
              emb_dim_per_rank_per_feature=self._emb_dim_per_rank_per_feature,
              device=self._device,
              callbacks=None,
              codecs=self._codecs,
          )
      else:
          self._dist = PooledEmbeddingsAllToAll(
              pg=self._pg,
              dim_sum_per_rank=self._dim_sum_per_rank,
              device=self._device,
              callbacks=self._callbacks,
              codecs=self._codecs,
          )
  
imports: from typing import Any, Callable, cast, Dict, List, Optional, TypeVar, Union
  import torch
  import torch.distributed
  from torchrec.distributed.dist_data import EmbeddingsAllToOne, KJTAllToAll, KJTOneToAll, PooledEmbeddingsAllToAll, VariableBatchPooledEmbeddingsAllToAll
  from torchrec.distributed.embedding_lookup import GroupedPooledEmbeddingsLookup, InferGroupedPooledEmbeddingsLookup
  from torchrec.distributed.embedding_sharding import BaseEmbeddingDist, BaseEmbeddingLookup, BaseSparseFeaturesDist, 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 3 for calling the TwPooledEmbeddingDist obj's underlying member method code with example usage: obj._create_output_dist_module(sharding_ctx) without return types