class_docstr: Bucketizes sparse features in TWRW 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.
      id_list_features_per_rank (List[int]): number of id list features to send to
          each rank.
      id_score_list_features_per_rank (List[int]): number of id score list features to
          send to each rank.
      id_list_feature_hash_sizes (List[int]): hash sizes of id list features.
      id_score_list_feature_hash_sizes (List[int]): hash sizes of id score list
          features.
      device (Optional[torch.device]): device on which buffers will be allocated.
      has_feature_processor (bool): existence of a feature processor (ie. position
          weighted features).
  
  Example::
  
      3 features
      2 hosts with 2 devices each
  
      Bucketize each feature into 2 buckets
      Staggered shuffle with feature splits [2, 1]
      AlltoAll operation
  
      NOTE: result of staggered shuffle and AlltoAll operation look the same after
      reordering in AlltoAll
  
      Result:
          host 0 device 0:
              feature 0 bucket 0
              feature 1 bucket 0
  
          host 0 device 1:
              feature 0 bucket 1
              feature 1 bucket 1
  
          host 1 device 0:
              feature 2 bucket 0
  
          host 1 device 1:
              feature 2 bucket 1
class_docstr_tok: ["Bucketizes","sparse","features","in","TWRW","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",".","id_list_features_per_rank","(","List","[","int","]",")",":","number","of","id","list","features","to","send","to","each","rank",".","id_score_list_features_per_rank","(","List","[","int","]",")",":","number","of","id","score","list","features","to","send","to","each","rank",".","id_list_feature_hash_sizes","(","List","[","int","]",")",":","hash","sizes","of","id","list","features",".","id_score_list_feature_hash_sizes","(","List","[","int","]",")",":","hash","sizes","of","id","score","list","features",".","device","(","Optional","[","torch.device","]",")",":","device","on","which","buffers","will","be","allocated",".","has_feature_processor","(","bool",")",":","existence","of","a","feature","processor","(","ie",".","position","weighted","features",")",".","Example",":",":","3","features","2","hosts","with","2","devices","each","Bucketize","each","feature","into","2","buckets","Staggered","shuffle","with","feature","splits","[","2",",","1","]","AlltoAll","operation","NOTE",":","result","of","staggered","shuffle","and","AlltoAll","operation","look","the","same","after","reordering","in","AlltoAll","Result",":","host","0","device","0",":","feature","0","bucket","0","feature","1","bucket","0","host","0","device","1",":","feature","0","bucket","1","feature","1","bucket","1","host","1","device","0",":","feature","2","bucket","0","host","1","device","1",":","feature","2","bucket","1"]
docstr: Reorders sparse data such that data is in contiguous blocks and correctly
  ordered for global TWRW layout.
code: def _staggered_shuffle(
      self, features_per_rank: List[int]
  ) -> List[int]:
      """
      Reorders sparse data such that data is in contiguous blocks and correctly
      ordered for global TWRW layout.
      """
  
      nodes = self._world_size // self._local_size
      features_per_node = [
          features_per_rank[node * self._local_size]
          for node in range(nodes)
      ]
      node_offsets = [0] + list(itertools.accumulate(features_per_node))
      num_features = node_offsets[-1]
  
      return [
          bucket * num_features + feature
          for node in range(nodes)
          for bucket in range(self._local_size)
          for feature in range(
              node_offsets[node], node_offsets[node + 1]
          )
      ]
  
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 3 for calling the TwRwSparseFeaturesDist obj's underlying member method code with example usage: obj._staggered_shuffle(features_per_rank) and returns: unknown