Add jaxtyping, fix minor bugs thus uncovered

curriculum_learning.py

@@ -110,6 +110,7 @@ def __init__(
         dist_backend: str = "nccl",
         local_rank: int = int(os.environ.get("LOCAL_RANK", 0)),
         llm_id: str = None,
+        smoke_test: bool = False,
     ):
         """
         Initialize the curriculum trainer.
@@ -132,6 +133,10 @@ def __init__(
         self.llm_id = llm_id
         self.llm_id_safe = self._sanitize_llm_id(llm_id)
 
+        # Smoke-test mode: 1 epoch, 1 train batch, 1 val batch, 1 eval batch.
+        # Exercises the full pipeline end-to-end for testing without a real run.
+        self.smoke_test = smoke_test
+
         # Distributed training parameters
         self.gradient_checkpointing = gradient_checkpointing
         self.dist_url = dist_url
@@ -784,6 +789,9 @@ def _evaluate_stage(
                             os.fsync(results_fp.fileno())
                         except Exception:
                             pass
+
+                    if self.smoke_test:
+                        break
         finally:
             if results_fp is not None:
                 results_fp.close()
@@ -903,6 +911,11 @@ def _train_stage(
         if batch_size is None:
             batch_size = BATCH_SIZE
 
+        # Smoke-test mode caps the run to a single epoch; the train/val/eval
+        # loops below additionally stop after one batch each.
+        if self.smoke_test:
+            num_epochs = 1
+
         if self.rank == 0:
             print(f"\n🚀 Starting {stage_name} Training with {self.model_type}")
             if eval_only:
@@ -1090,6 +1103,7 @@ def _train_stage(
                 prog = tqdm(
                     train_loader,
                     desc=f"Epoch {epoch}/{num_epochs}",
+                    total=1 if self.smoke_test else None,
                     disable=self.rank != 0,
                 )
                 for i, batch in enumerate(prog):
@@ -1128,12 +1142,17 @@ def _train_stage(
                             lr=f"{scheduler.get_last_lr()[0]:.2e}",
                         )
 
-                avg_train_loss = running_loss / len(train_loader)
+                    if self.smoke_test:
+                        break
+
+                num_train_batches = i + 1
+                avg_train_loss = running_loss / num_train_batches
                 if self.rank == 0:
                     tqdm.write(f"Epoch {epoch} — train loss: {avg_train_loss:.4f}")
 
                 # Validation
                 val_loss = 0.0
+                num_val_batches = 0
                 self.model.eval()
                 with torch.no_grad():
                     for batch in tqdm(
@@ -1142,8 +1161,11 @@ def _train_stage(
                         disable=self.rank != 0,
                     ):
                         val_loss += self._get_model().compute_loss(batch).item()
+                        num_val_batches += 1
+                        if self.smoke_test:
+                            break
 
-                avg_val_loss = val_loss / len(val_loader)
+                avg_val_loss = val_loss / num_val_batches
 
                 # Synchronize validation loss across all ranks
                 if dist.is_initialized():
@@ -1682,6 +1704,14 @@ def main():
         help="Local GPU rank",
     )
 
+    parser.add_argument(
+        "--smoke_test",
+        default=False,
+        action="store_true",
+        help="Fast smoke run: 1 epoch, 1 training batch, 1 validation batch, and 1 "
+        "evaluation batch, then done. Exercises the full pipeline end-to-end for testing.",
+    )
+
     # Logging arguments
     parser.add_argument(
         "--verbose", default=False, action="store_true", help="Enable verbose logging"
@@ -1702,6 +1732,7 @@ def main():
         dist_backend=args.dist_backend,
         local_rank=args.local_rank,
         llm_id=args.llm_id,
+        smoke_test=args.smoke_test,
     )
 
     # Run curriculum

pyproject.toml

@@ -55,6 +55,7 @@ dependencies = [
     "einops>=0.6",
     "wfdb>=4.0",
     "open-flamingo>=0.0.2",
+    "jaxtyping>=0.2.20",
 ]
 
 [project.urls]
@@ -64,6 +65,7 @@ Repository = "https://github.com/StanfordBDHG/OpenTSLM"
 
 [dependency-groups]
 dev = [
+    "beartype>=0.22.9",
     "reuse>=6.2.0",
     "ruff>=0.14.0",
 ]
@@ -100,3 +102,9 @@ members = [
 [tool.ruff]
 line-length = 120
 target-version = "py312"
+
+[tool.ruff.lint]
+# jaxtyping shape annotations like Float[Tensor, "batch length"] use string
+# literals inside the subscript, which pyflakes misreads as malformed forward
+# references (F722). This is the standard jaxtyping + ruff configuration.
+ignore = ["F722"]

requirements.txt

@@ -17,4 +17,5 @@ requests>=2.32.5
 einops>=0.8.1
 wfdb>=4.3.0
 open-flamingo>=0.0.2
+jaxtyping>=0.2.20
 -e .

src/opentslm/__init__.py

@@ -3,6 +3,21 @@
 #
 # SPDX-License-Identifier: MIT
 
+import os
+
+# Runtime type/shape checking is opt-in via OPENTSLM_MODE, evaluated once at import.
+#   "opt" (default): jaxtyping annotations stay inert -> zero overhead, and
+#                    beartype is not required to import the package.
+#   "dev":           install jaxtyping's beartype import hook so that Float[...]
+#                    shape annotations (and ordinary type hints) across the whole
+#                    opentslm package are enforced at runtime.
+# The hook must be installed before any opentslm submodule is imported, so this
+# block stays at the top of the package __init__, ahead of the imports below.
+if os.environ.get("OPENTSLM_MODE", "opt").lower() == "dev":
+    from jaxtyping import install_import_hook
+
+    install_import_hook("opentslm", "beartype.beartype")
+
 from opentslm.model.llm.OpenTSLM import OpenTSLM
 
 __all__ = ["OpenTSLM"]

src/opentslm/model/encoder/CNNTokenizer.py

@@ -5,6 +5,7 @@
 
 import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 
 from opentslm.model_config import TRANSFORMER_INPUT_DIM, ENCODER_OUTPUT_DIM, PATCH_SIZE
@@ -51,7 +52,9 @@ def __init__(
         self.input_norm = nn.LayerNorm(transformer_input_dim)
         self.input_dropout = nn.Dropout(self.dropout)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: Float[torch.Tensor, "batch length"]
+    ) -> Float[torch.Tensor, "batch patches embed_dim"]:
         """
         Args:
             x: FloatTensor of shape [B, L], a batch of raw time series.

src/opentslm/model/encoder/TimeSeriesEncoderBase.py

@@ -6,6 +6,7 @@
 from abc import abstractmethod
 import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 from opentslm.model_config import ENCODER_OUTPUT_DIM
 
@@ -21,5 +22,7 @@ def __init__(
         self.dropout = dropout
 
     @abstractmethod
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: Float[torch.Tensor, "batch length"]
+    ) -> Float[torch.Tensor, "batch patches embed_dim"]:
         pass

src/opentslm/model/encoder/TransformerCNNEncoder.py

@@ -5,6 +5,7 @@
 
 import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 
 from opentslm.model_config import TRANSFORMER_INPUT_DIM, ENCODER_OUTPUT_DIM, PATCH_SIZE
@@ -65,7 +66,9 @@ def __init__(
         )
         self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: Float[torch.Tensor, "batch length"]
+    ) -> Float[torch.Tensor, "batch patches embed_dim"]:
         """
         Args:
             x: FloatTensor of shape [B, L], a batch of raw time series.

src/opentslm/model/encoder/TransformerMLPEncoder.py

@@ -5,6 +5,7 @@
 
 import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 
 from opentslm.model_config import TRANSFORMER_INPUT_DIM, ENCODER_OUTPUT_DIM, PATCH_SIZE
@@ -33,7 +34,8 @@ def __init__(
             dropout: dropout probability
             max_patches: maximum number of patches expected per sequence (for pos emb)
         """
-        super().__init__(input_dim, output_dim, dropout)
+        super().__init__(output_dim, dropout)
+        self.input_dim = input_dim
         self.patch_size = patch_size
 
         if input_dim % patch_size != 0:
@@ -60,7 +62,9 @@ def __init__(
         )
         self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: Float[torch.Tensor, "batch length"]
+    ) -> Float[torch.Tensor, "batch patches embed_dim"]:
         """
         Args:
             x: FloatTensor of shape [B, L], a batch of raw time series.

src/opentslm/model/llm/OpenTSLMFlamingo.py

@@ -12,7 +12,8 @@
 from open_flamingo.src.utils import extend_instance
 import torch
 import torch._dynamo
-from typing import List, Dict, Tuple
+from typing import Any, List, Dict, Optional, Tuple
+from jaxtyping import Float, Int
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from opentslm.model_config import ENCODER_OUTPUT_DIM
@@ -153,9 +154,16 @@ def _infer_decoder_layers_attr_name(model):
         self.text_tokenizer = text_tokenizer
 
     def pad_and_apply_batch(
-        self, batch: List[Dict[str, any]], include_labels: bool
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        def pad_time_series(batch, max_length=None):
+        self, batch: List[Dict[str, Any]], include_labels: bool
+    ) -> Tuple[
+        Int[torch.Tensor, "batch seq_len"],
+        Float[torch.Tensor, "batch n_series 1 length"],
+        Int[torch.Tensor, "batch seq_len"],
+        Optional[Int[torch.Tensor, "batch seq_len"]],
+    ]:
+        def pad_time_series(
+            batch, max_length=None
+        ) -> Float[torch.Tensor, "batch n_series length"]:
             """Pad time series to the same length (either max in batch or specified max)"""
             time_series = [item["time_series"] for item in batch]
 
@@ -190,11 +198,16 @@ def pad_time_series(batch, max_length=None):
             "input_ids"
         ][-1]
 
-        # Process time series data
+        # Process time series data -> [B, n_series, length]
         images = pad_time_series(batch).to(
             self.device, dtype=cast_dtype, non_blocking=True
         )
-        images = images.unsqueeze(1)  # Add time dimension
+        # Flamingo expects vision input of shape [B, T_media, F_frames, ...].
+        # Each series must be its own media chunk (T_media == n_series) rather
+        # than separate frames of a single chunk, because Flamingo's masked
+        # cross-attention aligns each text token with the media chunk whose
+        # index equals the running count of preceding <image> tokens.
+        images = images.unsqueeze(2)
 
         # Process text inputs WITH answers
         text_inputs = []
@@ -244,7 +257,7 @@ def pad_time_series(batch, max_length=None):
         return input_ids, images, attention_mask, labels
 
     def generate(
-        self, batch: List[Dict[str, any]], max_new_tokens: int = 50, **generate_kwargs
+        self, batch: List[Dict[str, Any]], max_new_tokens: int = 50, **generate_kwargs
     ) -> List[str]:
         # Temporarily disable compilation to avoid data-dependent operation issues
         original_disable = torch._dynamo.config.disable
@@ -276,7 +289,7 @@ def generate(
             # Restore original compilation setting
             torch._dynamo.config.disable = original_disable
 
-    def compute_loss(self, batch: List[Dict[str, any]]) -> torch.Tensor:
+    def compute_loss(self, batch: List[Dict[str, Any]]) -> Float[torch.Tensor, ""]:
         """
         batch: same format as generate()
         answers: List[str] of length B

src/opentslm/model/llm/OpenTSLMSP.py

@@ -5,7 +5,8 @@
 
 import torch
 import torch.nn as nn
-from typing import List, Dict, Tuple, Optional
+from typing import Any, List, Dict, Tuple, Optional
+from jaxtyping import Float, Int
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from torch.nn.utils.rnn import pad_sequence
 
@@ -177,8 +178,11 @@ def disable_lora(self):
 
     def pad_and_apply_batch(
         self,
-        batch: List[Dict[str, any]],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch: List[Dict[str, Any]],
+    ) -> Tuple[
+        Float[torch.Tensor, "batch seq_len llm_dim"],
+        Int[torch.Tensor, "batch seq_len"],
+    ]:
         """
         TL;DR:
             This function is probably the most crucial part of OpenTSLM-SP, and also the hardest to understand.
@@ -316,7 +320,7 @@ def pad_and_apply_batch(
         return inputs_embeds, attention_mask
 
     def generate(
-        self, batch: List[Dict[str, any]], max_new_tokens: int = 50, **generate_kwargs
+        self, batch: List[Dict[str, Any]], max_new_tokens: int = 50, **generate_kwargs
     ) -> List[str]:
         inputs_embeds, attention_mask = self.pad_and_apply_batch(batch)
         gen_ids = self.llm.generate(
@@ -327,7 +331,7 @@ def generate(
         )
         return self.tokenizer.batch_decode(gen_ids, skip_special_tokens=True)
 
-    def compute_loss(self, batch: List[Dict[str, any]]) -> torch.Tensor:
+    def compute_loss(self, batch: List[Dict[str, Any]]) -> Float[torch.Tensor, ""]:
         """
         batch: same format as generate()
         answers: List[str] of length B

src/opentslm/model/llm/TimeSeriesFlamingoWithTrainableEncoder.py

@@ -5,6 +5,7 @@
 
 import torch
 from torch import nn
+from jaxtyping import Float
 from open_flamingo import Flamingo
 from einops import rearrange
 
@@ -25,7 +26,11 @@ def __init__(
     # Override the _encode_vision_x method to handle time series data
     # In the original Flamingo, the vision_encoder is a CLIPModel, which is not trainable (with torch.no_grad())
     # Here, we use a TimeSeriesCNNEncoder, which is trainable
-    def _encode_vision_x(self, vision_x):
+    def _encode_vision_x(
+        self,
+        vision_x: Float[torch.Tensor, "batch t_media n_frames length"]
+        | Float[torch.Tensor, "batch t_media n_frames channels height width"],
+    ) -> None:
         # Handle time series data while still using the TimeSeriesCNNEncoder
         if vision_x.ndim == 4:  # For shape (b, T_img, F, features)
             b, T, F, features = vision_x.shape

src/opentslm/model/llm/TimeSeriesLLM.py

@@ -7,6 +7,7 @@
 
 import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 from opentslm.prompt.full_prompt import FullPrompt
 
@@ -25,7 +26,7 @@ def generate(
 
         raise NotImplementedError("Generate method should be implemented by the subclass")
 
-    def compute_loss(self, batch: List[Dict[str, Any]]) -> torch.Tensor:
+    def compute_loss(self, batch: List[Dict[str, Any]]) -> Float[torch.Tensor, ""]:
         """
         batch: same format as generate()
         answers: List[str] of length B