code: def multi_head_attention(
      query,
      key,
      value,
      num_heads,
      q_weight,
      k_weight,
      v_weight,
      out_weight,
      q_bias=None,
      k_bias=None,
      v_bias=None,
      out_bias=None,
      attn_bias_k=None,
      attn_bias_v=None,
      dropout=0.0,
      additive_mask=None,
      key_padding_mask=None,
  ):
      """MultiHeadAttention.
  
      Computes multi-headed attention with query, key, and value.
      We use the following notations to describe the inputs and outputs below.
      :math:`L_T`: target sequence length, :math:`L_S`: source sequence length, :math:`B`: batch size, :math:`D`: input dimension, :math:`E`: embedding dimension, :math:`H`: number of attention heads.
  
      References:
  
          A. Vaswani et al. "Attention is All You Need."
          NIPS. 2017.
          <https://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf>
  
      Args:
          query (~nnabla.Variable): Input N-D array with shape :math:`(L_T, B, D_q)`.
          key (~nnabla.Variable): Input N-D array with shape :math:`(L_S, B, D_k)`.
          value (~nnabla.Variable): Input N-D array with shape :math:`(L_S, B, D_v)`.
          num_heads (int): Number of attention heads. Note that embedding dimensoin E must be divisible by the number of heads. Default is 12 which is conventional.
          q_weight (~nnabla.Variable): Input N-D array with shape :math:`(D_q, E)`.
          k_weight (~nnabla.Variable): Input N-D array with shape :math:`(D_k, E)`.
          v_weight (~nnabla.Variable): Input N-D array with shape :math:`(D_v, E_v)`.
          out_weight (~nnabla.Variable): Input N-D array with shape :math:`(D_v, E_{out})`.
          q_bias (~nnabla.Variable, optional): Input N-D array with shape :math:`(E, )`.
          k_bias (~nnabla.Variable, optional): Input N-D array with shape :math:`(E, )`.
          v_bias (~nnabla.Variable, optional): Input N-D array with shape :math:`(E_v, )`.
          out_bias (~nnabla.Variable, optional): Input N-D array with shape :math:`(E_{out}, )`.
          attn_bias_k (~nnabla.Variable, optional): Input N-D array with shape :math:`(E, )`.
          attn_bias_v (~nnabla.Variable, optional): Input N-D array with shape :math:`(E_v, )`.
          dropout (float, optional): Dropout ratio applied to parameters. Default is 0.
          additive_mask (~nnabla.Variable, optional): Input N-D array with shape :math:`(L_T, L_S)`. Values will be added to the attention layer to prevent attention to certain positions.
          key_padding_mask (~nnabla.Variable, optional): Input N-D array with shape :math:`(B, L_S)`. Specified padding elements will be ignored by the attention layer. Values must be either 1 or 0.
  
      Returns:
          ~nnabla.Variable: Output :math:`y` with shape :math:`(L_T, B, E_{out})`
          ~nnabla.Variable: Output :math:`h_n` with shape :math:`(B, L_T, L_S)`
      """
  
      from . import functions as F
  
      tgt_len, batch_size, _ = query.shape
      src_len, batch_size, _ = key.shape
      q_embed_dim = q_weight.shape[1]
      k_embed_dim = k_weight.shape[1]
      v_embed_dim = v_weight.shape[1]
      out_dim = out_weight.shape[1]
      assert src_len == value.shape[0]
      head_dim = q_embed_dim // num_heads
      head_vdim = v_embed_dim // num_heads
      assert (
          q_embed_dim == k_embed_dim
      ), "embedding dimensions must be the same for query and key."
      assert head_dim * num_heads == q_embed_dim, (
          "embedding dimension must be divisibile by num_heads %d"
          % num_heads
      )
      assert head_vdim * num_heads == v_embed_dim, (
          "v_embed_dim must be divisibile by num_heads %d." % num_heads
      )
  
      if key_padding_mask is not None:
          assert key_padding_mask.shape[0] == batch_size
          assert key_padding_mask.shape[1] == src_len
  
      # query:(L_T, B, E) --> q:(L_T, B, E)
      q = F.affine(query, q_weight, q_bias, base_axis=2)
      # key:(L_S, B, D_k) --> k:(L_S, B, E_k)
      k = F.affine(key, k_weight, k_bias, base_axis=2)
      # value:(L_S, B, D_v) --> v:(L_S, B, E_v)
      v = F.affine(value, v_weight, v_bias, base_axis=2)
  
      q *= float(head_dim) ** -0.5
  
      if attn_bias_k is not None:
          attn_bias_k = F.reshape(attn_bias_k, (1, 1, k_embed_dim))
          attn_bias_v = F.reshape(attn_bias_v, (1, 1, v_embed_dim))
          src_len += 1
          assert attn_bias_k is not None
          attn_bias_k = F.broadcast(
              attn_bias_k, (1, batch_size, attn_bias_k.shape[2])
          )
          attn_bias_v = F.broadcast(
              attn_bias_v, (1, batch_size, attn_bias_v.shape[2])
          )
          k = F.concatenate(k, attn_bias_k, axis=0)
          v = F.concatenate(v, attn_bias_v, axis=0)
          if additive_mask is not None:
              # additive_mask: (L_T, L_S) --> (L_T, L_S + 1)
              additive_mask = F.pad(additive_mask, (0, 1))
          if key_padding_mask is not None:
              # key_padding_mask: (B, L_S) --> (B, L_S + 1)
              key_padding_mask = F.pad(key_padding_mask, (0, 1))
  
      q = F.transpose(
          F.reshape(q, (tgt_len, batch_size * num_heads, head_dim)),
          (1, 0, 2),
      )  # q:(B*H, L_T, head_dim)
      k = F.transpose(
          F.reshape(k, (-1, batch_size * num_heads, head_dim)),
          (1, 0, 2),
      )  # k:(B*H, L_S, head_dim)
      v = F.transpose(
          F.reshape(v, (-1, batch_size * num_heads, head_vdim)),
          (1, 0, 2),
      )  # v:(B*H, L_S, head_vdim)
  
      # attn_output_weights: (B*H, L_T, L_S)
      attn_output_weights = F.batch_matmul(q, k, transpose_b=True)
      assert list(attn_output_weights.shape) == [
          batch_size * num_heads,
          tgt_len,
          src_len,
      ]
  
      if additive_mask is not None:
          additive_mask = F.reshape(
              additive_mask, ((1,) + additive_mask.shape)
          )
          attn_output_weights += additive_mask
  
      if key_padding_mask is not None:
          attn_output_weights = F.reshape(
              attn_output_weights,
              (batch_size, num_heads, tgt_len, src_len),
          )
          attn_output_weights = F.where(
              F.broadcast(
                  F.reshape(
                      key_padding_mask, (batch_size, 1, 1, src_len)
                  ),
                  attn_output_weights.shape,
              ),  # Condition
              F.constant(
                  val=float("-inf"), shape=attn_output_weights.shape
              ),  # If true
              attn_output_weights,
          )  # If false
          attn_output_weights = F.reshape(
              attn_output_weights,
              (batch_size * num_heads, tgt_len, src_len),
          )
  
      attn_output_weights = F.softmax(
          attn_output_weights, axis=len(attn_output_weights.shape) - 1
      )
      if dropout > 0:
          attn_output_weights = F.dropout(
              attn_output_weights, p=dropout
          )
  
      # (B*H, L_T, L_S) x (B*H, L_S, head_vdim) --> (B*H, L_T, head_vdim)
      attn_output = F.batch_matmul(attn_output_weights, v)
      assert list(attn_output.shape) == [
          batch_size * num_heads,
          tgt_len,
          head_vdim,
      ]
      attn_output = F.reshape(
          F.transpose(attn_output, (1, 0, 2)),
          (tgt_len, batch_size, v_embed_dim),
      )  # attn_output: (L_T, B, E_v)
  
      attn_output = F.affine(
          attn_output, out_weight, out_bias, base_axis=2
      )
  
      return attn_output, attn_output_weights
  
imports: from __future__ import absolute_import
  from .function_bases import *
  from .None import dtypes
  import nnabla
  import numpy
  from .normalization_functions import *
  from .numpy_compat_functions import *
label_desc: Use image node_id 29 for calling a global function with example usage: multi_head_attention(query, key, value, num_heads, q_weight, k_weight, v_weight, out_weight, q_bias, k_bias, v_bias, out_bias, attn_bias_k, attn_bias_v, dropout, additive_mask, key_padding_mask) and returns: attn_output,  attn_output_weights