def add_decomposed_rel_pos(
    attn: Tensor, q: Tensor, rel_pos_h: Tensor, rel_pos_w: Tensor, q_size: tuple[int, int], k_size: tuple[int, int]
) -> Tensor:
    """Calculate decomposed Relative Positional Embeddings.

    from :paper:`mvitv2`.
    https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py


    Args:
        attn: attention map.
        q: query q in the attention layer with shape (B, q_h * q_w, C).
        rel_pos_h: relative position embeddings (Lh, C) for height axis.
        rel_pos_w: relative position embeddings (Lw, C) for width axis.
        q_size: spatial sequence size of query q with (q_h, q_w).
        k_size: spatial sequence size of key k with (k_h, k_w).

    Returns:
        att: attention map with added relative positional embeddings.

    """
    q_h, q_w = q_size
    k_h, k_w = k_size
    Rh = get_rel_pos(q_h, k_h, rel_pos_h)
    Rw = get_rel_pos(q_w, k_w, rel_pos_w)

    B, _, dim = q.shape
    r_q = q.reshape(B, q_h, q_w, dim)
    rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)
    rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)

    attn = (attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]).view(
        B, q_h * q_w, k_h * k_w
    )

    return attn
Calculate decomposed Relative Positional Embeddings.

    from :paper:`mvitv2`.
    https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py


    Args:
        attn: attention map.
        q: query q in the attention layer with shape (B, q_h * q_w, C).
        rel_pos_h: relative position embeddings (Lh, C) for height axis.
        rel_pos_w: relative position embeddings (Lw, C) for width axis.
        q_size: spatial sequence size of query q with (q_h, q_w).
        k_size: spatial sequence size of key k with (k_h, k_w).

    Returns:
        att: attention map with added relative positional embeddings.

    
kornia/contrib/models/sam/architecture/image_encoder.py