lgatr.primitives.attention

Equivariant attention.

Functions

scaled_dot_product_attention(query, key, ...)

Execute scaled dot-product attention.

sdp_attention(q_mv, k_mv, v_mv[, q_s, k_s, v_s])

Equivariant geometric attention based on scaled dot products.
lgatr.primitives.attention.scaled_dot_product_attention(query, key, value, **attn_kwargs)[source]

Execute scaled dot-product attention.

The attention backend is determined dynamically based on the attn_kwargs provided (see lgatr.primitives.attention_backends.get_attention_backend()).

Parameters:

  • query (Tensor) – Tensor of shape (..., items_out, channels).

  • key (Tensor) – Tensor of shape (..., items_in, channels).

  • value (Tensor) – Tensor of shape (..., items_in, channels).

  • **attn_kwargs – Optional keyword arguments forwarded to the attention backend.

Returns:

Tensor of shape (..., items_out, channels).

Return type:

outputs

lgatr.primitives.attention.sdp_attention(q_mv, k_mv, v_mv, q_s=None, k_s=None, v_s=None, **attn_kwargs)[source]

Equivariant geometric attention based on scaled dot products.

Expects multivector and (optionally) scalar queries, keys, and values, and computes:

attn_weights[..., i, j] = softmax_j[
    ga_inner_product(q_mv[..., i, :, :], k_mv[..., j, :, :])
    + euclidean_inner_product(q_s[..., i, :], k_s[..., j, :])
]
outputs_mv[..., i, c, :] = sum_j attn_weights[..., i, j] v_mv[..., j, c, :]
outputs_s[..., i, c]     = sum_j attn_weights[..., i, j] v_s[..., j, c]
Parameters:

  • q_mv (Tensor) – Multivector queries of shape (..., items_out, mv_channels, 16).

  • k_mv (Tensor) – Multivector keys of shape (..., items_in, mv_channels, 16).

  • v_mv (Tensor) – Multivector values of shape (..., items_in, mv_channels, 16).

  • q_s (Tensor | None) – Optional scalar queries of shape (..., items_out, s_channels). If None, the scalar inner product is omitted and outputs_s is None.

  • k_s (Tensor | None) – Optional scalar keys of shape (..., items_in, s_channels). Must be None iff q_s is None.

  • v_s (Tensor | None) – Optional scalar values of shape (..., items_in, s_channels). Must be None iff q_s is None.

  • **attn_kwargs – Optional keyword arguments forwarded to the attention backend.

Return type:

tuple[Tensor, Tensor | None]

Returns:

  • outputs_mv – Multivector result of shape (..., items_out, mv_channels, 16).

  • outputs_s – Scalar result of shape (..., items_out, s_channels), or None if q_s is None.