class_docstr: Multihead Attention with extra flags on the q/k/v and out projections.
class_docstr_tok: ["Multihead","Attention","with","extra","flags","on","the","q\/k\/v","and","out","projections","."]
code: def forward(
      self,
      query,
      key,
      value,
      key_padding_mask=None,
      need_weights=False,
      attn_mask=None,
      do_qkv_proj=True,
      do_out_proj=True,
      rpe=True,
  ):
      tgt_len, bsz, embed_dim = query.shape
      head_dim = embed_dim // self.num_heads
      v_head_dim = self.vdim // self.num_heads
      assert (
          head_dim * self.num_heads == embed_dim
      ), "embed_dim must be divisible by num_heads"
      scaling = float(head_dim) ** -0.5
  
      # whether or not use the original query/key/value
      q = self.q_proj(query) * scaling if do_qkv_proj else query
      k = self.k_proj(key) if do_qkv_proj else key
      v = self.v_proj(value) if do_qkv_proj else value
  
      if attn_mask is not None:
          dtype_lst = [
              paddle.float32,
              paddle.float64,
              paddle.float16,
              paddle.uint8,
              paddle.bool,
          ]
          assert (
              attn_mask.dtype in dtype_lst
          ), f"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}"
          if attn_mask.dtype == paddle.uint8:
              msg = "Byte tensor for attn_mask in nn.MultiHeadAttention is deprecated. Use bool tensor instead."
              logger.warning(msg)
              attn_mask = attn_mask.to(paddle.bool)
  
          if attn_mask.dim() == 2:
              attn_mask = attn_mask.unsqueeze(0)
              if list(attn_mask.shape) != [
                  1,
                  query.shape[0],
                  key.shape[0],
              ]:
                  raise RuntimeError(
                      "The size of the 2D attn_mask is not correct."
                  )
          elif attn_mask.dim() == 3:
              if attn_mask.shape != [
                  bsz * self.num_heads,
                  query.shape[0],
                  key.shape[0],
              ]:
                  raise RuntimeError(
                      "The size of the 3D attn_mask is not correct."
                  )
          else:
              raise RuntimeError(
                  f"attn_mask's dimension {attn_mask.dim()} is not supported"
              )
  
      # convert ByteTensor key_padding_mask to bool
      if (
          key_padding_mask is not None
          and key_padding_mask.dtype == paddle.uint8
      ):
          msg = "Byte tensor for key_padding_mask in nn.MultiHeadAttention is deprecated. Use bool tensor instead."
          logger.warning(msg)
          key_padding_mask = key_padding_mask.to(paddle.bool)
  
      q = q.reshape(
          [tgt_len, bsz * self.num_heads, head_dim]
      ).transpose([1, 0, 2])
      if k is not None:
          k = k.reshape([-1, bsz * self.num_heads, head_dim]).transpose(
              [1, 0, 2]
          )
      if v is not None:
          v = v.reshape(
              [-1, bsz * self.num_heads, v_head_dim]
          ).transpose([1, 0, 2])
  
      src_len = k.shape[1]
  
      if key_padding_mask is not None:
          assert key_padding_mask.shape[0] == bsz
          assert key_padding_mask.shape[1] == src_len
  
      if self.add_zero_attn:
          src_len += 1
          k = paddle.concat(
              [
                  k,
                  paddle.zeros(
                      (k.shape[0], 1) + k.shape[2:], dtype=k.dtype
                  ),
              ],
              axis=1,
          )
          v = paddle.concat(
              [
                  v,
                  paddle.zeros(
                      (v.shape[0], 1) + v.shape[2:], dtype=v.dtype
                  ),
              ],
              axis=1,
          )
  
          if attn_mask is not None:
              attn_mask = F.pad(attn_mask, (0, 1))
          if key_padding_mask is not None:
              key_padding_mask = F.pad(key_padding_mask, (0, 1))
  
      attn_output_weights = paddle.bmm(q, k.transpose([0, 2, 1]))
      assert list(attn_output_weights.shape) == [
          bsz * self.num_heads,
          tgt_len,
          src_len,
      ]
      """ Add relative position embedding."""
      if self.rpe and rpe:
          # NOTE: for simplicity, we assume src_len == tgt_len == window_size**2 here
          assert (
              src_len == self.window_size[0] * self.window_size[1]
              and tgt_len == self.window_size[0] * self.window_size[1]
          ), f"src{src_len}, tgt{tgt_len}, window{self.window_size[0]}"
          relative_position_bias = self.relative_position_bias_table[
              self.relative_position_index.reshape([-1])
          ].reshape(
              [
                  self.window_size[0] * self.window_size[1],
                  self.window_size[0] * self.window_size[1],
                  -1,
              ]
          )  # Wh*Ww,Wh*Ww,nH
          relative_position_bias = relative_position_bias.transpose(
              [2, 0, 1]
          )  # nH, Wh*Ww, Wh*Ww
          attn_output_weights = attn_output_weights.reshape(
              [bsz, self.num_heads, tgt_len, src_len]
          ) + relative_position_bias.unsqueeze(0)
          attn_output_weights = attn_output_weights.reshape(
              [bsz * self.num_heads, tgt_len, src_len]
          )
      # Attention weight for the invalid region is -inf.
      if attn_mask is not None:
          if attn_mask.dtype == paddle.bool:
              attn_output_weights.masked_fill_(attn_mask, float("-inf"))
          else:
              attn_output_weights += attn_mask
  
      if key_padding_mask is not None:
          attn_output_weights = attn_output_weights.reshape(
              [bsz, self.num_heads, tgt_len, src_len]
          )
          attn_output_weights = attn_output_weights.masked_fill(
              key_padding_mask.unsqueeze(1).unsqueeze(2),
              float("-inf"),
          )
          attn_output_weights = attn_output_weights.reshape(
              [bsz * self.num_heads, tgt_len, src_len]
          )
  
      attn_output_weights = F.softmax(attn_output_weights, axis=-1)
      attn_output_weights = F.dropout(
          attn_output_weights, p=self.dropout, training=self.training
      )
  
      attn_output = paddle.bmm(attn_output_weights, v)
      assert list(attn_output.shape) == [
          bsz * self.num_heads,
          tgt_len,
          v_head_dim,
      ]
      attn_output = attn_output.transpose([1, 0, 2]).reshape(
          [tgt_len, bsz, self.vdim]
      )
      if do_out_proj:
          attn_output = F.linear(
              attn_output, self.out_proj.weight, self.out_proj.bias
          )
  
      if need_weights:
          # average attention weights over heads
          attn_output_weights = attn_output_weights.reshape(
              [bsz, self.num_heads, tgt_len, src_len]
          )
          return (
              attn_output,
              q,
              k,
              attn_output_weights.sum(axis=1) / self.num_heads,
          )
      else:
          return attn_output, q, k
  
imports: import paddle
  import paddle.nn
  import paddle.nn.functional
  from paddleseg.cvlibs import manager, param_init
  from paddleseg.models import layers
  from paddleseg.utils import utils, logger
  from paddleseg.models.backbones.transformer_utils import *
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(query, key, value, key_padding_mask, need_weights, attn_mask, do_qkv_proj, do_out_proj, rpe) and returns: attn_output,  q,  k,  unknown, attn_output,  q,  k