tensorplay.nn.attention.create_mask

tensorplay.nn.attention.create_mask(mod_fn: _score_mod_signature | _mask_mod_signature, B: int | None, H: int | None, Q_LEN: int, KV_LEN: int, device: DeviceLikeType | None = None) → Tensor[source]

This function creates a mask tensor from a mod_fn function.

Parameters:

• mod_fn (Union[_score_mod_signature, _mask_mod_signature]) – Function to modify attention scores.

• B (int) – Batch size.

• H (int) – Number of query heads.

• Q_LEN (int) – Sequence length of query.

• KV_LEN (int) – Sequence length of key/value.

• device (str) – Device to run the mask creation on.

Returns:

A mask tensor with shape (B, H, M, N).

Return type:

mask (Tensor)