class_docstr: Fixed (3 line):
      assert embed_dim == self.embed_dim -> remove
      assert list(query.size()) == [tgt_len, bsz, embed_dim] -> remove
      attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim) -> attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim)
docstr: Input shape: Time(SeqLen) x Batch x Channel
docstr_tok: ["Input","shape",":","Time","(","SeqLen",")","x","Batch","x","Channel"]
code: def forward(
      self,
      query,
      key: Optional[Tensor],
      key_padding_mask: Optional[Tensor] = None,
      layer_state: Optional[Dict[str, Optional[Tensor]]] = None,
      attn_mask: Optional[Tensor] = None,
      layer_head_mask: Optional[Tensor] = None,
      output_attentions=False,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      """Input shape: Time(SeqLen) x Batch x Channel"""
      static_kv: bool = self.encoder_decoder_attention
      tgt_len, bsz, embed_dim = query.size()
      # get here for encoder decoder cause of static_kv
      if layer_state is not None:  # reuse k,v and encoder_padding_mask
          saved_state = layer_state.get(self.cache_key, {})
          if "prev_key" in saved_state and static_kv:
              # previous time steps are cached - no need to recompute key and value if they are static
              key = None
      else:
          saved_state = None
          layer_state = {}
  
      q = self.q_proj(query) * self.scaling
      if static_kv:
          if key is None:
              k = v = None
          else:
              k = self.k_proj(key)
              v = self.v_proj(key)
      else:
          k = self.k_proj(query)
          v = self.v_proj(query)
  
      q = self._shape(q, tgt_len, bsz)
      if k is not None:
          k = self._shape(k, -1, bsz)
      if v is not None:
          v = self._shape(v, -1, bsz)
  
      if saved_state is not None:
          k, v, key_padding_mask = self._use_saved_state(
              k, v, saved_state, key_padding_mask, static_kv, bsz
          )
  
      # Update cache
      layer_state[self.cache_key] = {
          "prev_key": k.view(bsz, self.num_heads, -1, self.head_dim),
          "prev_value": v.view(bsz, self.num_heads, -1, self.head_dim),
          "prev_key_padding_mask": key_padding_mask
          if not static_kv
          else None,
      }
  
      assert k is not None
      src_len = k.size(1)
      attn_weights = torch.bmm(q, k.transpose(1, 2))
      assert attn_weights.size() == (
          bsz * self.num_heads,
          tgt_len,
          src_len,
      )
  
      if attn_mask is not None:
          attn_weights = (
              attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              + attn_mask
          )
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      # This is part of a workaround to get around fork/join parallelism not supporting Optional types.
      if key_padding_mask is not None and key_padding_mask.dim() == 0:
          key_padding_mask = None
      assert key_padding_mask is None or key_padding_mask.size()[
          :2
      ] == (
          bsz,
          src_len,
      )
  
      if (
          key_padding_mask is not None
      ):  # don't attend to padding symbols
          attn_weights = attn_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          reshaped = key_padding_mask.unsqueeze(1).unsqueeze(2)
          attn_weights = attn_weights.masked_fill(
              reshaped, float("-inf")
          )
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      attn_weights = nn.functional.softmax(attn_weights, dim=-1)
  
      if layer_head_mask is not None:
          assert layer_head_mask.size() == (
              self.num_heads,
          ), f"Head mask for a single layer should be of size {(self.num_heads,)}, but is {layer_head_mask.size()}"
          attn_weights = layer_head_mask.view(
              1, -1, 1, 1
          ) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
          attn_weights = attn_weights.view(
              bsz * self.num_heads, tgt_len, src_len
          )
  
      if output_attentions:
          # make sure that attn_weights are included in graph
          attn_weights_reshaped = attn_weights.view(
              bsz, self.num_heads, tgt_len, src_len
          )
          attn_weights = attn_weights_reshaped.view(
              bsz * self.num_heads, tgt_len, src_len
          )
      else:
          attn_weights_reshaped = None
  
      attn_probs = nn.functional.dropout(
          attn_weights,
          p=self.dropout,
          training=self.training,
      )
  
      assert v is not None
      attn_output = torch.bmm(attn_probs, v)
      assert attn_output.size() == (
          bsz * self.num_heads,
          tgt_len,
          self.head_dim,
      )
      attn_output = (
          attn_output.transpose(0, 1)
          .contiguous()
          .view(tgt_len, bsz, self.embed_dim)
      )
      attn_output = self.out_proj(attn_output)
  
      return attn_output, attn_weights_reshaped
  
imports: from typing import Dict, Optional, Tuple
  import torch
  from torch import Tensor, nn
  from transformers.models.fsmt.modeling_fsmt import Attention
label_desc: Use image node_id 1 for calling the Attention_ obj's underlying member method code with example usage: obj.forward(query, key, key_padding_mask, layer_state, attn_mask, layer_head_mask, output_attentions) and returns: attn_output,  attn_weights_reshaped