tensorplay.nn.functional.multi_head_attention_forward

tensorplay.nn.functional.multi_head_attention_forward(query: TensorBase, key: TensorBase, value: TensorBase, embed_dim_to_check: int, num_heads: int, in_proj_weight: TensorBase | None = None, in_proj_bias: TensorBase | None = None, bias_k=None, bias_v=None, add_zero_attn: bool = False, dropout_p: float = 0.0, out_proj_weight: TensorBase | None = None, out_proj_bias: TensorBase | None = 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]

torch-compatible multi_head_attention_forward.

Follows the structure of third_party/pytorch/aten/src/ATen/native/transformers/attention.cpp _scaled_dot_product_attention paths composed from dispatched primitives.