class_docstr: Stores KJTAllToAll context and reuses it in SequenceEmbeddingsAllToAll.
  SequenceEmbeddingsAllToAll has the same comm pattern as KJTAllToAll.
  
  Attributes:
      features_before_input_dist (Optional[KeyedJaggedTensor]): stores the original
          KJT before input dist.
      input_splits(List[int]): stores the input splits of KJT AlltoAll.
      output_splits (List[int]): stores the output splits of KJT AlltoAll.
      unbucketize_permute_tensor (Optional[torch.Tensor]): stores the permute order of
          KJT bucketize (for row-wise sharding only).
      lengths_after_input_dist (Optional[torch.Tensor]): stores the KJT length after
          input dist.
class_docstr_tok: ["Stores","KJTAllToAll","context","and","reuses","it","in","SequenceEmbeddingsAllToAll",".","SequenceEmbeddingsAllToAll","has","the","same","comm","pattern","as","KJTAllToAll",".","Attributes",":","features_before_input_dist","(","Optional","[","KeyedJaggedTensor","]",")",":","stores","the","original","KJT","before","input","dist",".","input_splits","(","List","[","int","]",")",":","stores","the","input","splits","of","KJT","AlltoAll",".","output_splits","(","List","[","int","]",")",":","stores","the","output","splits","of","KJT","AlltoAll",".","unbucketize_permute_tensor","(","Optional","[","torch.Tensor","]",")",":","stores","the","permute","order","of","KJT","bucketize","(","for","row-wise","sharding","only",")",".","lengths_after_input_dist","(","Optional","[","torch.Tensor","]",")",":","stores","the","KJT","length","after","input","dist","."]
code: def record_stream(self, stream: torch.cuda.streams.Stream) -> None:
      if self.features_before_input_dist is not None:
          self.features_before_input_dist.record_stream(stream)
      if self.sparse_features_recat is not None:
          self.sparse_features_recat.record_stream(stream)
      if self.unbucketize_permute_tensor is not None:
          self.unbucketize_permute_tensor.record_stream(stream)
      if self.lengths_after_input_dist is not None:
          self.lengths_after_input_dist.record_stream(stream)
  
file_path: torchrec/torchrec/distributed/sharding/sequence_sharding.py
imports: from dataclasses import dataclass, field
  from typing import List, Optional
  import torch
  from torchrec.distributed.embedding_sharding import EmbeddingShardingContext
  from torchrec.distributed.embedding_types import KJTList
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.streamable import Multistreamable
label_desc: Use image node_id 1 for calling the SequenceShardingContext obj's underlying member method code with example usage: obj.record_stream(stream) without return types