Type valid token masks explicitly

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 8a255805-976d-4261-b430-561adafb3cc6

src/colipri/pooling.py

@@ -5,12 +5,13 @@ from torch import nn
 from .types import TypePooledEmbeddings
 from .types import TypeSequenceEmbeddings
 from .types import TypeTextAttentionMask
+from .types import TypeValidTokenMask
 
 
 def _get_valid_token_mask(
     sequence: TypeSequenceEmbeddings,
     attention_mask: TypeTextAttentionMask | None,
-) -> torch.Tensor | None:
+) -> TypeValidTokenMask | None:
     """Get a validated Boolean mask for valid sequence positions.
 
     Args:
@@ -44,7 +45,7 @@ def _get_valid_token_mask(
 
 def _compute_masked_mean(
     sequence: TypeSequenceEmbeddings,
-    valid_mask: torch.Tensor | None,
+    valid_mask: TypeValidTokenMask | None,
 ) -> TypePooledEmbeddings:
     """Compute the sequence mean while excluding masked positions."""
     if valid_mask is None:

src/colipri/types.py

@@ -2,6 +2,7 @@ import os
 from typing import TypeAlias
 
 import torchio as tio
+from jaxtyping import Bool
 from jaxtyping import Float
 from jaxtyping import Int64
 from torch import Tensor
@@ -25,6 +26,7 @@ TypeImagesTensor: TypeAlias = Float[Tensor, "batch 1 x_in y_in z_in"]
 
 TypeTextTokenIds: TypeAlias = Int64[Tensor, "batch text_length"]
 TypeTextAttentionMask: TypeAlias = Int64[Tensor, "batch text_length"]
+TypeValidTokenMask: TypeAlias = Bool[Tensor, "batch sequence_length"]
 
 # Outputs
 TypeSequenceEmbeddings: TypeAlias = Float[Tensor, "batch sequence_length embed_dim"]