tensorplay.nn.functional.multi_head_attention_forward

tensorplay.nn.functional.multi_head_attention_forward(query: Tensor, key: Tensor, value: Tensor, embed_dim_to_check: int, num_heads: int, in_proj_weight: Tensor = None, in_proj_bias: Tensor = None, bias_k=None, bias_v=None, add_zero_attn: bool = False, dropout_p: float = 0.0, out_proj_weight: Tensor = None, out_proj_bias: Tensor = None, training: bool = True, key_padding_mask=None, need_weights: bool = True, attn_mask=None, use_separate_proj_weight: bool = False, q_proj_weight=None, k_proj_weight=None, v_proj_weight=None, static_k=None, static_v=None, average_attn_weights: bool = True, is_causal: bool = False)[source]

multi_head_attention_forward(query, key, value, embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias=None, bias_k=None, bias_v=None, add_zero_attn=False, dropout_p=0.0, out_proj_weight=None, out_proj_bias=None, training=True, key_padding_mask=None, need_weights=True, attn_mask=None, use_separate_proj_weight=False, q_proj_weight=None, k_proj_weight=None, v_proj_weight=None, static_k=None, static_v=None, average_attn_weights=True, is_causal=False) -> (Tensor, Optional[Tensor])

Computes multi-head attention on (L, N, E) inputs (2D unbatched inputs are promoted internally and the batch dim is squeezed on return). Returns the projected output of shape (L, N, E) and, when need_weights is true, the attention weights of shape (N, L, S) — or (num_heads, L, S) with average_attn_weights=False.