class_docstr: Initialization.
  
  Arguments:
      local_attention (Module): local attention with q,k,v
      sequence_process_group (ProcessGroup): sequence parallel process group
      scatter_idx (int): scatter_idx for all2all comm
      gather_idx (int): gather_idx for all2all comm
class_docstr_tok: ["Initialization",".","Arguments",":","local_attention","(","Module",")",":","local","attention","with","q",",","k",",","v","sequence_process_group","(","ProcessGroup",")",":","sequence","parallel","process","group","scatter_idx","(","int",")",":","scatter_idx","for","all2all","comm","gather_idx","(","int",")",":","gather_idx","for","all2all","comm"]
code: def __init__(
      self,
      local_attention: Module,
      sequence_process_group: dist.ProcessGroup,
      scatter_idx: int = 2,
      gather_idx: int = 0,
  ) -> None:
      super(DistributedAttention, self).__init__()
      self.local_attn = local_attention
      self.spg = sequence_process_group
      self.scatter_idx = scatter_idx
      self.gather_idx = gather_idx
  
file_path: DeepSpeed/deepspeed/sequence/layer.py
imports: import torch
  from typing import Any, Tuple
  from torch import Tensor
  from torch.nn import Module
  import deepspeed.comm
label_desc: Use image node_id 1 to create a new DistributedAttention object from inherited base classes: torch.nn with example: obj = DistributedAttention(local_attention, sequence_process_group, scatter_idx, gather_idx)