class_docstr: Redistributes sparse features with an AlltoAll collective operation for table wise
  sharding.
  
  Args:
      pg (dist.ProcessGroup): ProcessGroup for AlltoAll communication.
      features_per_rank (List[int]): number of features to send to each rank.
class_docstr_tok: ["Redistributes","sparse","features","with","an","AlltoAll","collective","operation","for","table","wise","sharding",".","Args",":","pg","(","dist.ProcessGroup",")",":","ProcessGroup","for","AlltoAll","communication",".","features_per_rank","(","List","[","int","]",")",":","number","of","features","to","send","to","each","rank","."]
code: def __init__(
      self,
      pg: dist.ProcessGroup,
      features_per_rank: List[int],
  ) -> None:
      super().__init__()
      self._dist = KJTAllToAll(
          pg=pg,
          splits=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 to create a new TwSparseFeaturesDist object from inherited base classes: BaseSparseFeaturesDist with example: obj = TwSparseFeaturesDist(pg, features_per_rank)