class_docstr: Shards embedding bags table-wise then row-wise.
class_docstr_tok: ["Shards","embedding","bags","table-wise","then","row-wise","."]
code: def create_input_dist(
      self, device: Optional[torch.device] = None
  ) -> BaseSparseFeaturesDist[KeyedJaggedTensor]:
      features_per_rank = self._features_per_rank(
          self._grouped_embedding_configs_per_rank
      )
      feature_hash_sizes = self._get_feature_hash_sizes()
      assert self._pg is not None
      assert self._intra_pg is not None
      return TwRwSparseFeaturesDist(
          pg=self._pg,
          local_size=self._intra_pg.size(),
          features_per_rank=features_per_rank,
          feature_hash_sizes=feature_hash_sizes,
          device=device if device is not None else self._device,
          has_feature_processor=self._has_feature_processor,
          need_pos=self._need_pos,
      )
  
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 for calling the TwRwPooledEmbeddingSharding obj's underlying member method code with example usage: obj.create_input_dist(device) and returns: TwRwSparseFeaturesDist