mask_to_kv_lengths

cuequivariance_torch.mask_to_kv_lengths(mask)

Convert a right-padded triangle-attention mask to per-row key/value lengths.

Parameters:

mask (torch.Tensor) – Boolean-like mask of shape (B, N, 1, 1, K). For B=1, can also be (N, 1, 1, K). The last dimension must be prefix-shaped: all True entries first, followed by all False entries.

Returns:

int32 tensor of shape (B, N, 1, 1, 1) containing each row’s effective key/value length. Zero-length rows are allowed.

Return type:

torch.Tensor

Raises:

ValueError – If the mask shape is not supported or the mask is not prefix-shaped.

Note

This helper validates the prefix contract. For torch.compile-heavy code, compute lengths before the compiled region and pass them to triangle_attention() with kv_lengths=....