lloca.backbone.attention_backends.mask.get_sparse_attention_mask

lloca.backbone.attention_backends.mask.get_sparse_attention_mask(batch, attention_backend, dtype)

Returns sparse attention mask according to the backend.

Parameters:

• batch (torch.Tensor) – Batch vector, maps each token to its sequence in the batch.

• attention_backend (str) – Attention backend to use (“varlen”, “xformers”, “flex”, or “flash”).

• dtype (torch.dtype) – Data type of the attention mask (for xformers backend).

Returns:

Attention mask for the specified backend.

Return type:

dict[str, torch.Tensor | BlockMask | BlockDiagonalMask]