class_docstr: Implements an efficient Sparse Self Attention of Transformer layer based on `Generative Modeling with Sparse Transformers`: https://arxiv.org/abs/1904.10509
  
  For more information please see, TODO DeepSpeed Sparse Transformer.
  
  For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial.
docstr: Applies forward phase of sparse self attention
  
  Arguments:
      query: required: query tensor
      key: required: key tensor
      value: required: value tensor
      rpe: optional: a tensor same dimension as x that is used as relative position embedding
      key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength)
      attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
      key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied
      attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied
  
  Return:
       attn_output: a dense tensor containing attention context
code: def forward(
      self,
      query,
      key,
      value,
      rpe=None,
      key_padding_mask=None,
      attn_mask=None,
  ):
      """Applies forward phase of sparse self attention
  
      Arguments:
          query: required: query tensor
          key: required: key tensor
          value: required: value tensor
          rpe: optional: a tensor same dimension as x that is used as relative position embedding
          key_padding_mask: optional: a mask tensor of size (BatchSize X SequenceLength)
          attn_mask: optional: a mask tensor of size (SequenceLength X SequenceLength); currently only 2D is supported
          key_padding_mask_mode: optional: a boolean determining if key_padding_mask needs to be added or multiplied
          attn_mask_mode: optional: a boolean determining if attn_mask needs to be added or multiplied
  
      Return:
           attn_output: a dense tensor containing attention context
      """
      assert (
          query.dtype == torch.half
      ), "sparse attention only supports training in fp16 currently, please file a github issue if you need fp32 support"
      bsz, num_heads, tgt_len, head_dim = query.size()
  
      # transpose back key if it is already transposed
      key = self.transpose_key_for_scores(key, tgt_len)
  
      # check that operation is supported
      if query.shape != key.shape or key.shape != value.shape:
          raise NotImplementedError(
              "only self-attention is supported for now"
          )
  
      # squeeze key_padding_mask if it is given
      if key_padding_mask is not None:
          key_padding_mask = self.transpose_mask_for_sparse(
              query.dtype, key_padding_mask, is_key_padding_mask=True
          )
  
      # squeeze attn_mask if it is given
      if attn_mask is not None:
          attn_mask = self.transpose_mask_for_sparse(
              query.dtype, attn_mask
          )
  
      # cache look-up table computations etc
      (
          sparse_dot_sdd_nt,
          sparse_dot_dsd_nn,
          sparse_softmax,
      ) = self.get_ops(num_heads, tgt_len)
  
      scaling = float(head_dim) ** -0.5
  
      # attention scores
      attn_output_weights = sparse_dot_sdd_nt(query, key)
      attn_output_weights = sparse_softmax(
          attn_output_weights,
          scale=scaling,
          rpe=rpe,
          key_padding_mask=key_padding_mask,
          attn_mask=attn_mask,
          key_padding_mask_mode=self.key_padding_mask_mode,
          attn_mask_mode=self.attn_mask_mode,
      )
  
      # outputs
      attn_output = sparse_dot_dsd_nn(attn_output_weights, value)
      return attn_output
  
file_path: DeepSpeed/deepspeed/ops/sparse_attention/sparse_self_attention.py
imports: import torch.nn
  import torch
  from torch import distributed
  from deepspeed.ops.sparse_attention import SparsityConfig
label_desc: Use image node_id 6 for calling the SparseSelfAttention obj's underlying member method code with example usage: obj.forward(query, key, value, rpe, key_padding_mask, attn_mask) and returns: attn_output