code: def multi_head_attention_forward(
      query: Tensor,
      key: Tensor,
      value: Tensor,
      embed_dim_to_check: int,
      num_heads: int,
      in_proj_weight: Tensor,
      in_proj_bias: Tensor,
      bias_k: Optional[Tensor],
      bias_v: Optional[Tensor],
      add_zero_attn: bool,
      dropout_p: float,
      out_proj_weight: Tensor,
      out_proj_bias: Tensor,
      training: bool = True,
      key_padding_mask: Optional[Tensor] = None,
      need_weights: bool = True,
      attn_mask: Optional[Tensor] = None,
      use_separate_proj_weight: bool = False,
      q_proj_weight: Optional[Tensor] = None,
      k_proj_weight: Optional[Tensor] = None,
      v_proj_weight: Optional[Tensor] = None,
      static_k: Optional[Tensor] = None,
      static_v: Optional[Tensor] = None,
      attention_probs_forward_hook=None,
      attention_probs_backwards_hook=None,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      if not torch.jit.is_scripting():
          tens_ops = (
              query,
              key,
              value,
              in_proj_weight,
              in_proj_bias,
              bias_k,
              bias_v,
              out_proj_weight,
              out_proj_bias,
          )
          if any(
              not isinstance(t, Tensor) for t in tens_ops
          ) and F.has_torch_function(tens_ops):
              return F.handle_torch_function(
                  multi_head_attention_forward,
                  tens_ops,
                  query,
                  key,
                  value,
                  embed_dim_to_check,
                  num_heads,
                  in_proj_weight,
                  in_proj_bias,
                  bias_k,
                  bias_v,
                  add_zero_attn,
                  dropout_p,
                  out_proj_weight,
                  out_proj_bias,
                  training=training,
                  key_padding_mask=key_padding_mask,
                  need_weights=need_weights,
                  attn_mask=attn_mask,
                  use_separate_proj_weight=use_separate_proj_weight,
                  q_proj_weight=q_proj_weight,
                  k_proj_weight=k_proj_weight,
                  v_proj_weight=v_proj_weight,
                  static_k=static_k,
                  static_v=static_v,
              )
      tgt_len, bsz, embed_dim = query.size()
      assert embed_dim == embed_dim_to_check
      # allow MHA to have different sizes for the feature dimension
      assert key.size(0) == value.size(0) and key.size(1) == value.size(
          1
      )
  
      head_dim = embed_dim // num_heads
      assert (
          head_dim * num_heads == embed_dim
      ), "embed_dim must be divisible by num_heads"
      scaling = float(head_dim) ** -0.5
  
      if not use_separate_proj_weight:
          if torch.equal(query, key) and torch.equal(key, value):
              # self-attention
              q, k, v = F.linear(
                  query, in_proj_weight, in_proj_bias
              ).chunk(3, dim=-1)
  
          elif torch.equal(key, value):
              # encoder-decoder attention
              # This is inline in_proj function with in_proj_weight and in_proj_bias
              b_ = in_proj_bias
              start_ = 0
              end_ = embed_dim
              w_ = in_proj_weight[start_:end_, :]
              if b_ is not None:
                  b_ = b_[start_:end_]
              q = F.linear(query, w_, b_)
  
              if key is None:
                  assert value is None
                  k = None
                  v = None
              else:
                  # This is inline in_proj function with in_proj_weight and in_proj_bias
                  b_ = in_proj_bias
                  start_ = embed_dim
                  end_ = None
                  w_ = in_proj_weight[start_:, :]
                  if b_ is not None:
                      b_ = b_[start_:]
                  k, v = F.linear(key, w_, b_).chunk(2, dim=-1)
  
          else:
              # This is inline in_proj function with in_proj_weight and in_proj_bias
              b_ = in_proj_bias
              start_ = 0
              end_ = embed_dim
              w_ = in_proj_weight[start_:end_, :]
              if b_ is not None:
                  b_ = b_[start_:end_]
              q = F.linear(query, w_, b_)
  
              # This is inline in_proj function with in_proj_weight and in_proj_bias
              b_ = in_proj_bias
              start_ = embed_dim
              end_ = embed_dim * 2
              w_ = in_proj_weight[start_:end_, :]
              if b_ is not None:
                  b_ = b_[start_:end_]
              k = F.linear(key, w_, b_)
  
              # This is inline in_proj function with in_proj_weight and in_proj_bias
              b_ = in_proj_bias
              start_ = embed_dim * 2
              end_ = None
              w_ = in_proj_weight[start_:, :]
              if b_ is not None:
                  b_ = b_[start_:]
              v = F.linear(value, w_, b_)
      else:
          q_proj_weight_non_opt = torch.jit._unwrap_optional(
              q_proj_weight
          )  # pylint: disable=protected-access
          len1, len2 = q_proj_weight_non_opt.size()
          assert len1 == embed_dim and len2 == query.size(-1)
  
          k_proj_weight_non_opt = torch.jit._unwrap_optional(
              k_proj_weight
          )  # pylint: disable=protected-access
          len1, len2 = k_proj_weight_non_opt.size()
          assert len1 == embed_dim and len2 == key.size(-1)
  
          v_proj_weight_non_opt = torch.jit._unwrap_optional(
              v_proj_weight
          )  # pylint: disable=protected-access
          len1, len2 = v_proj_weight_non_opt.size()
          assert len1 == embed_dim and len2 == value.size(-1)
  
          if in_proj_bias is not None:
              q = F.linear(
                  query,
                  q_proj_weight_non_opt,
                  in_proj_bias[0:embed_dim],
              )
              k = F.linear(
                  key,
                  k_proj_weight_non_opt,
                  in_proj_bias[embed_dim : (embed_dim * 2)],
              )
              v = F.linear(
                  value,
                  v_proj_weight_non_opt,
                  in_proj_bias[(embed_dim * 2) :],
              )
          else:
              q = F.linear(query, q_proj_weight_non_opt, in_proj_bias)
              k = F.linear(key, k_proj_weight_non_opt, in_proj_bias)
              v = F.linear(value, v_proj_weight_non_opt, in_proj_bias)
      q = q * scaling
  
      if attn_mask is not None:
          assert attn_mask.dtype in (
              torch.float32,
              torch.float64,
              torch.float16,
              torch.uint8,
              torch.bool,
          ), f"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}"
          if attn_mask.dtype == torch.uint8:
              warnings.warn(
                  "Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
              )
              attn_mask = attn_mask.to(torch.bool)
  
          if attn_mask.dim() == 2:
              attn_mask = attn_mask.unsqueeze(0)
              if list(attn_mask.size()) != [
                  1,
                  query.size(0),
                  key.size(0),
              ]:
                  raise RuntimeError(
                      "The size of the 2D attn_mask is not correct."
                  )
          elif attn_mask.dim() == 3:
              if list(attn_mask.size()) != [
                  bsz * num_heads,
                  query.size(0),
                  key.size(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"
              )
          # attn_mask"s dim is 3 now.
  
      # convert ByteTensor key_padding_mask to bool
      if (
          key_padding_mask is not None
          and key_padding_mask.dtype == torch.uint8
      ):
          warnings.warn(
              "Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead."
          )
          key_padding_mask = key_padding_mask.to(torch.bool)
  
      if bias_k is not None and bias_v is not None:
          if static_k is None and static_v is None:
              k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
              v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
              if attn_mask is not None:
                  attn_mask = pad(attn_mask, (0, 1))
              if key_padding_mask is not None:
                  key_padding_mask = pad(key_padding_mask, (0, 1))
          else:
              assert (
                  static_k is None
              ), "bias cannot be added to static key."
              assert (
                  static_v is None
              ), "bias cannot be added to static value."
      else:
          assert bias_k is None
          assert bias_v is None
  
      q = (
          q.contiguous()
          .view(tgt_len, bsz * num_heads, head_dim)
          .transpose(0, 1)
      )
      if k is not None:
          k = (
              k.contiguous()
              .view(-1, bsz * num_heads, head_dim)
              .transpose(0, 1)
          )
      if v is not None:
          v = (
              v.contiguous()
              .view(-1, bsz * num_heads, head_dim)
              .transpose(0, 1)
          )
  
      if static_k is not None:
          assert static_k.size(0) == bsz * num_heads
          assert static_k.size(2) == head_dim
          k = static_k
  
      if static_v is not None:
          assert static_v.size(0) == bsz * num_heads
          assert static_v.size(2) == head_dim
          v = static_v
  
      src_len = k.size(1)
  
      if key_padding_mask is not None:
          assert key_padding_mask.size(0) == bsz
          assert key_padding_mask.size(1) == src_len
  
      if add_zero_attn:
          src_len += 1
          k = torch.cat(
              [
                  k,
                  torch.zeros(
                      (k.size(0), 1) + k.size()[2:],
                      dtype=k.dtype,
                      device=k.device,
                  ),
              ],
              dim=1,
          )
          v = torch.cat(
              [
                  v,
                  torch.zeros(
                      (v.size(0), 1) + v.size()[2:],
                      dtype=v.dtype,
                      device=v.device,
                  ),
              ],
              dim=1,
          )
          if attn_mask is not None:
              attn_mask = pad(attn_mask, (0, 1))
          if key_padding_mask is not None:
              key_padding_mask = pad(key_padding_mask, (0, 1))
  
      attn_output_weights = torch.bmm(q, k.transpose(1, 2))
      assert list(attn_output_weights.size()) == [
          bsz * num_heads,
          tgt_len,
          src_len,
      ]
  
      if attn_mask is not None:
          if attn_mask.dtype == torch.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.view(
              bsz, 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.view(
              bsz * num_heads, tgt_len, src_len
          )
  
      attn_output_weights = F.softmax(attn_output_weights, dim=-1)
      attn_output_weights = F.dropout(
          attn_output_weights, p=dropout_p, training=training
      )
  
      # use hooks for the attention weights if necessary
      if (
          attention_probs_forward_hook is not None
          and attention_probs_backwards_hook is not None
      ):
          attention_probs_forward_hook(attn_output_weights)
          attn_output_weights.register_hook(
              attention_probs_backwards_hook
          )
  
      attn_output = torch.bmm(attn_output_weights, v)
      assert list(attn_output.size()) == [
          bsz * num_heads,
          tgt_len,
          head_dim,
      ]
      attn_output = (
          attn_output.transpose(0, 1)
          .contiguous()
          .view(tgt_len, bsz, embed_dim)
      )
      attn_output = F.linear(
          attn_output, out_proj_weight, out_proj_bias
      )
  
      if need_weights:
          # average attention weights over heads
          attn_output_weights = attn_output_weights.view(
              bsz, num_heads, tgt_len, src_len
          )
          return attn_output, attn_output_weights.sum(dim=1) / num_heads
      else:
          return attn_output, None
  
file_path: towhee/towhee/models/clip/auxilary.py
imports: import warnings
  from typing import Tuple, Optional
  import torch
  from torch import Tensor
  from torch.nn.init import xavier_uniform_
  from torch.nn.init import constant_
  from torch.nn.init import xavier_normal_
  from torch.nn.parameter import Parameter
  from torch.nn import functional
label_desc: Use image node_id 1 for calling a global function with example usage: multi_head_attention_forward(query, key, value, embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias, bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight, out_proj_bias, training, key_padding_mask, need_weights, attn_mask, use_separate_proj_weight, q_proj_weight, k_proj_weight, v_proj_weight, static_k, static_v, attention_probs_forward_hook, attention_probs_backwards_hook) and returns: attn_output,  unknown, attn_output,  None, F