class_docstr: Bucketizes sparse features in RW fashion and then redistributes with an AlltoAll
  collective operation.
  
  Args:
      pg (dist.ProcessGroup): ProcessGroup for AlltoAll communication.
      intra_pg (dist.ProcessGroup): ProcessGroup within single host group for AlltoAll
          communication.
      num_features (int): total number of features.
      feature_hash_sizes (List[int]): hash sizes of features.
      device (Optional[torch.device]): device on which buffers will be allocated.
      is_sequence (bool): if this is for a sequence embedding.
      has_feature_processor (bool): existence of feature processor (ie. position
          weighted features).
docstr: Bucketizes sparse feature values into world size number of buckets and then
  performs AlltoAll operation.
  
  Args:
      sparse_features (KeyedJaggedTensor): sparse features to bucketize and
          redistribute.
  
  Returns:
      Awaitable[Awaitable[KeyedJaggedTensor]]: awaitable of awaitable of KeyedJaggedTensor.
docstr_tok: ["Bucketizes","sparse","feature","values","into","world","size","number","of","buckets","and","then","performs","AlltoAll","operation",".","Args",":","sparse_features","(","KeyedJaggedTensor",")",":","sparse","features","to","bucketize","and","redistribute",".","Returns",":","Awaitable","[","Awaitable","[","KeyedJaggedTensor","]","]",":","awaitable","of","awaitable","of","KeyedJaggedTensor","."]
code: def forward(
      self,
      sparse_features: KeyedJaggedTensor,
  ) -> Awaitable[Awaitable[KeyedJaggedTensor]]:
      """
      Bucketizes sparse feature values into world size number of buckets and then
      performs AlltoAll operation.
  
      Args:
          sparse_features (KeyedJaggedTensor): sparse features to bucketize and
              redistribute.
  
      Returns:
          Awaitable[Awaitable[KeyedJaggedTensor]]: awaitable of awaitable of KeyedJaggedTensor.
      """
  
      (
          bucketized_features,
          self.unbucketize_permute_tensor,
      ) = bucketize_kjt_before_all2all(
          sparse_features,
          num_buckets=self._world_size,
          block_sizes=self._feature_block_sizes_tensor,
          output_permute=self._is_sequence,
          bucketize_pos=self._has_feature_processor
          if sparse_features.weights_or_none() is None
          else self._need_pos,
      )
  
      return self._dist(bucketized_features)
  
imports: import math
  from typing import Any, cast, Dict, List, Optional, Tuple, TypeVar, Union
  import torch
  import torch.distributed
  from torchrec.distributed.dist_data import EmbeddingsAllToOneReduce, KJTAllToAll, KJTOneToAll, PooledEmbeddingsReduceScatter, VariableBatchPooledEmbeddingsReduceScatter
  from torchrec.distributed.embedding_lookup import GroupedPooledEmbeddingsLookup, InferGroupedPooledEmbeddingsLookup
  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, 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 2 for calling the RwSparseFeaturesDist obj's underlying member method code with example usage: obj.forward(sparse_features) and returns: self