class_docstr: Shards embedding bags table-wise, i.e.. a given embedding table is entirely placed
  on a selected rank.
class_docstr_tok: ["Shards","embedding","bags","table-wise",",","i.e","..","a","given","embedding","table","is","entirely","placed","on","a","selected","rank","."]
code: def create_input_dist(
      self,
      device: Optional[torch.device] = None,
  ) -> BaseSparseFeaturesDist[KeyedJaggedTensor]:
      assert self._pg is not None
      return TwSparseFeaturesDist(
          self._pg,
          self.features_per_rank(),
      )
  
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 1 for calling the TwPooledEmbeddingSharding obj's underlying member method code with example usage: obj.create_input_dist(device) and returns: TwSparseFeaturesDist