class_docstr: Redistributes sequence embedding tensor in TW fashion with an AlltoAll operation.
  
  Args:
      pg (dist.ProcessGroup): ProcessGroup for AlltoAll communication.
      features_per_rank (List[int]): number of features (sum of dimensions) of the
          embedding for each rank.
      device (Optional[torch.device]): device on which buffers will be allocated.
class_docstr_tok: ["Redistributes","sequence","embedding","tensor","in","TW","fashion","with","an","AlltoAll","operation",".","Args",":","pg","(","dist.ProcessGroup",")",":","ProcessGroup","for","AlltoAll","communication",".","features_per_rank","(","List","[","int","]",")",":","number","of","features","(","sum","of","dimensions",")","of","the","embedding","for","each","rank",".","device","(","Optional","[","torch.device","]",")",":","device","on","which","buffers","will","be","allocated","."]
docstr: Performs AlltoAll operation on sequence embeddings tensor.
  
  Args:
      local_embs (torch.Tensor): tensor of values to distribute.
      sharding_ctx (SequenceShardingContext): shared context from KJTAllToAll
          operation.
  
  Returns:
      Awaitable[torch.Tensor]: awaitable of sequence embeddings.
docstr_tok: ["Performs","AlltoAll","operation","on","sequence","embeddings","tensor",".","Args",":","local_embs","(","torch.Tensor",")",":","tensor","of","values","to","distribute",".","sharding_ctx","(","SequenceShardingContext",")",":","shared","context","from","KJTAllToAll","operation",".","Returns",":","Awaitable","[","torch.Tensor","]",":","awaitable","of","sequence","embeddings","."]
code: def forward(
      self,
      local_embs: torch.Tensor,
      sharding_ctx: Optional[SequenceShardingContext] = None,
  ) -> Awaitable[torch.Tensor]:
      """
      Performs AlltoAll operation on sequence embeddings tensor.
  
      Args:
          local_embs (torch.Tensor): tensor of values to distribute.
          sharding_ctx (SequenceShardingContext): shared context from KJTAllToAll
              operation.
  
      Returns:
          Awaitable[torch.Tensor]: awaitable of sequence embeddings.
      """
  
      assert sharding_ctx is not None
      return self._dist(
          local_embs,
          lengths=sharding_ctx.lengths_after_input_dist,
          input_splits=sharding_ctx.input_splits,
          output_splits=sharding_ctx.output_splits,
          batch_size_per_rank=sharding_ctx.batch_size_per_rank,
          sparse_features_recat=sharding_ctx.sparse_features_recat,
          unbucketize_permute_tensor=None,
      )
  
imports: from typing import Any, Dict, List, Optional
  import torch
  import torch.distributed
  from torchrec.distributed.dist_data import SeqEmbeddingsAllToOne, SequenceEmbeddingsAllToAll
  from torchrec.distributed.embedding_lookup import GroupedEmbeddingsLookup, InferGroupedEmbeddingsLookup
  from torchrec.distributed.embedding_sharding import BaseEmbeddingDist, BaseEmbeddingLookup, BaseSparseFeaturesDist
  from torchrec.distributed.embedding_types import BaseGroupedFeatureProcessor, KJTList
  from torchrec.distributed.sharding.sequence_sharding import InferSequenceShardingContext, SequenceShardingContext
  from torchrec.distributed.sharding.tw_sharding import BaseTwEmbeddingSharding, InferTwSparseFeaturesDist, TwSparseFeaturesDist
  from torchrec.distributed.types import Awaitable, CommOp, QuantizedCommCodecs
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the TwSequenceEmbeddingDist obj's underlying member method code with example usage: obj.forward(local_embs, sharding_ctx) and returns: self