Function _merge_masks

imperative/python/megengine/functional/nn.py:2201–2323

r""" Determine mask type and combine masks if necessary. Note: This function will continue to improve with the iteration of MHA. Args: attn_mask: MHA's attention mask tensor, the shape is :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)` key_padding_mask: MHA's

(
    attn_mask: Tensor,
    key_padding_mask: Tensor,
    query: Tensor,
    key: Tensor,
    add_bias_kv: bool = False,
    add_zero_attn: bool = False,
    is_causal: bool = False,
    maybe_cudnn_style_mask: bool = False,
    num_heads: int = 0,
)

def _merge_masks(
    attn_mask: Tensor,
    key_padding_mask: Tensor,
    query: Tensor,
    key: Tensor,
    add_bias_kv: bool = False,
    add_zero_attn: bool = False,
    is_causal: bool = False,
    maybe_cudnn_style_mask: bool = False,
    num_heads: int = 0,
):
    r"""
    Determine mask type and combine masks if necessary.

    Note: This function will continue to improve with the iteration of MHA.

    Args:
        attn_mask: MHA's attention mask tensor, the shape is :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`
        key_padding_mask: MHA's padding mask tensor, the shape is :math:`(N, S)`
        query: MHA's query, the shape is :math:`(N, L, E_q)`
        key: MHA's key, the shape is :math:`(N, S, E_k)`
        add_bias_kv: used to determine whether pad is needed on the sequence dimension of attn_mask and key_padding_mask, from MHA's ``add_bias_kv``.
        add_zero_attn: used to determine whether pad is needed on the sequence dimension of attn_mask and key_padding_mask, from MHA's ``add_zero_attn``.
        is_causal: MHA's is_causal, is_causal provides a hint that attn_mask is the causal mask.
        maybe_cudnn_style_mask: MHA's maybe_cudnn_style_mask, like is_causal, maybe_cudnn_style_mask provides a hint that attn_mask and key_padding_mask is the cudnn style mask.
        num_heads: MHA's head number.
    Returns:
        merged_mask: merged mask, may be None, the shape is :math:`(L, S)`, :math:`(2\cdotL + 2\cdotN)` or :math:`(N\cdot\text{num\_heads}, L, S)`
        mask_type: merged mask type ``("no_mask", "default_mask", "cudnn_style_mask" or "user_defined_mask")``
    """
    mask_type = "no_mask"
    merged_mask = None
    batch_size = query.shape[0]
    seq_qlen = query.shape[1]
    seq_klen = key.shape[1]
    attn_mask_dim = attn_mask.ndim if attn_mask is not None else 0
    attn_mask_shape = attn_mask.shape if attn_mask is not None else 0
    key_padding_mask_shape = (
        key_padding_mask.shape if key_padding_mask is not None else 0
    )

    # is_causal is used to hint whether to use a causal mask, where the upper right triangle is all -inf,
    # and the diagonal and lower left triangle are all 0. But if attn_mask is given, attn_mask is used first.
    if (
        not maybe_cudnn_style_mask
        and is_causal
        and attn_mask is None
        and key_padding_mask is None
    ):
        # At this point, merged_mask = None
        mask_type = "default_mask"
    elif (
        not maybe_cudnn_style_mask
        and is_causal
        and attn_mask is not None
        and key_padding_mask is None
    ):
        # At this point, merged_mask = attn_mask