diff --git a/curriculum_learning.py b/curriculum_learning.py
index 4b25fe25..a7a49391 100644
--- a/curriculum_learning.py
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index 6f7056a0..54b92057 100644
--- a/pyproject.toml
+++ b/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"]
diff --git a/requirements.txt b/requirements.txt
index 097fbbe2..58d90f4b 100644
--- a/requirements.txt
+++ b/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 .
\ No newline at end of file
diff --git a/src/opentslm/__init__.py b/src/opentslm/__init__.py
index 0078c70d..a0591911 100644
--- a/src/opentslm/__init__.py
+++ b/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"]
\ No newline at end of file
diff --git a/src/opentslm/model/encoder/CNNTokenizer.py b/src/opentslm/model/encoder/CNNTokenizer.py
index be945cb0..2cef01eb 100644
--- a/src/opentslm/model/encoder/CNNTokenizer.py
+++ b/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.
diff --git a/src/opentslm/model/encoder/TimeSeriesEncoderBase.py b/src/opentslm/model/encoder/TimeSeriesEncoderBase.py
index 99204152..7ea9db0f 100644
--- a/src/opentslm/model/encoder/TimeSeriesEncoderBase.py
+++ b/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
diff --git a/src/opentslm/model/encoder/TransformerCNNEncoder.py b/src/opentslm/model/encoder/TransformerCNNEncoder.py
index 6cd01b04..aaf51e20 100644
--- a/src/opentslm/model/encoder/TransformerCNNEncoder.py
+++ b/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.
diff --git a/src/opentslm/model/encoder/TransformerMLPEncoder.py b/src/opentslm/model/encoder/TransformerMLPEncoder.py
index 89403ea3..c86ec79d 100644
--- a/src/opentslm/model/encoder/TransformerMLPEncoder.py
+++ b/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.
diff --git a/src/opentslm/model/llm/OpenTSLMFlamingo.py b/src/opentslm/model/llm/OpenTSLMFlamingo.py
index b0f5676b..ef31072b 100644
--- a/src/opentslm/model/llm/OpenTSLMFlamingo.py
+++ b/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
diff --git a/src/opentslm/model/llm/OpenTSLMSP.py b/src/opentslm/model/llm/OpenTSLMSP.py
index 4e7735d8..69247c50 100644
--- a/src/opentslm/model/llm/OpenTSLMSP.py
+++ 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
diff --git a/src/opentslm/model/llm/TimeSeriesFlamingoWithTrainableEncoder.py b/src/opentslm/model/llm/TimeSeriesFlamingoWithTrainableEncoder.py
index 97a7b52e..a8c36a36 100644
--- a/src/opentslm/model/llm/TimeSeriesFlamingoWithTrainableEncoder.py
+++ 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
diff --git a/src/opentslm/model/llm/TimeSeriesLLM.py b/src/opentslm/model/llm/TimeSeriesLLM.py
index b3a568d5..4c91dc23 100644
--- a/src/opentslm/model/llm/TimeSeriesLLM.py
+++ b/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
diff --git a/src/opentslm/model/projector/LinearProjector.py b/src/opentslm/model/projector/LinearProjector.py
index 0dfab497..a9ea0bab 100644
--- a/src/opentslm/model/projector/LinearProjector.py
+++ b/src/opentslm/model/projector/LinearProjector.py
@@ -3,7 +3,9 @@
 #
 # SPDX-License-Identifier: MIT
 
+import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 
 class LinearProjector(nn.Module):
@@ -11,5 +13,7 @@ def __init__(self, input_dim, output_dim, device):
         super().__init__()
         self.projector = nn.Linear(input_dim, output_dim).to(device)
 
-    def forward(self, x):
+    def forward(
+        self, x: Float[torch.Tensor, "*batch patches encoder_dim"]
+    ) -> Float[torch.Tensor, "*batch patches llm_dim"]:
         return self.projector(x)
diff --git a/src/opentslm/model/projector/MLPProjector.py b/src/opentslm/model/projector/MLPProjector.py
index e8c9a3a1..fd0728fc 100644
--- a/src/opentslm/model/projector/MLPProjector.py
+++ b/src/opentslm/model/projector/MLPProjector.py
@@ -3,7 +3,9 @@
 #
 # SPDX-License-Identifier: MIT
 
+import torch
 import torch.nn as nn
+from jaxtyping import Float
 
 
 class MLPProjector(nn.Module):
@@ -17,5 +19,7 @@ def __init__(self, input_dim, output_dim, device):
             nn.Dropout(0.0),
         ).to(device)
 
-    def forward(self, x):
+    def forward(
+        self, x: Float[torch.Tensor, "*batch patches encoder_dim"]
+    ) -> Float[torch.Tensor, "*batch patches llm_dim"]:
         return self.projector(x)
diff --git a/src/opentslm/prompt/text_time_series_prompt.py b/src/opentslm/prompt/text_time_series_prompt.py
index bdcbadd8..0ffe0492 100644
--- a/src/opentslm/prompt/text_time_series_prompt.py
+++ b/src/opentslm/prompt/text_time_series_prompt.py
@@ -6,6 +6,7 @@
 from .prompt import Prompt
 import numpy as np
 from collections.abc import Sequence
+from jaxtyping import Num
 
 
 class TextTimeSeriesPrompt(Prompt):
@@ -39,5 +40,5 @@ def __init__(self, text: str, time_series: Sequence):
     def get_text(self) -> str:
         return self.__text
 
-    def get_time_series(self) -> np.ndarray:
+    def get_time_series(self) -> Num[np.ndarray, " length"]:
         return self.__time_series
diff --git a/src/opentslm/time_series_datasets/monash/monash_utils.py b/src/opentslm/time_series_datasets/monash/monash_utils.py
index cc6a5520..539d8ae0 100644
--- a/src/opentslm/time_series_datasets/monash/monash_utils.py
+++ b/src/opentslm/time_series_datasets/monash/monash_utils.py
@@ -5,6 +5,7 @@
 
 import numpy as np
 import pandas as pd
+from jaxtyping import Shaped
 from tqdm import tqdm
 
 import requests
@@ -87,7 +88,7 @@ def load_from_tsfile_to_dataframe(
     full_file_path_and_name,
     return_separate_X_and_y=True,
     replace_missing_vals_with="NaN",
-):
+) -> tuple[pd.DataFrame, Shaped[np.ndarray, " n_cases"]] | pd.DataFrame:
     """Loads data from a .ts file into a Pandas DataFrame.
 
     Parameters
diff --git a/src/opentslm/time_series_datasets/ucr/ucr_loader.py b/src/opentslm/time_series_datasets/ucr/ucr_loader.py
index 0ccf44d9..05657b82 100644
--- a/src/opentslm/time_series_datasets/ucr/ucr_loader.py
+++ b/src/opentslm/time_series_datasets/ucr/ucr_loader.py
@@ -10,6 +10,7 @@
 
 import pandas as pd
 import torch
+from jaxtyping import Float, Int
 from torch.utils.data import Dataset, DataLoader
 
 # from constants import RAW_DATA_PATH
@@ -120,7 +121,7 @@ def __init__(
     def __len__(self):
         return len(self.df)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx) -> tuple[Float[torch.Tensor, " length"], int]:
         row    = self.df.iloc[idx]
         feats  = row[self.feature_cols].astype(float).values
         tensor = torch.tensor(feats, dtype=torch.float32)
@@ -129,7 +130,9 @@ def __getitem__(self, idx):
         label  = int(row[self.label_col])
         return tensor, label
 
-def collate_fn(batch):
+def collate_fn(
+    batch: list[tuple[Float[torch.Tensor, " length"], int]],
+) -> tuple[Float[torch.Tensor, "batch length"], Int[torch.Tensor, " batch"]]:
     """
     Stack into:
       - features: FloatTensor (batch_size, series_length)
diff --git a/src/opentslm/time_series_datasets/util.py b/src/opentslm/time_series_datasets/util.py
index 8b924db2..1e3aef7e 100644
--- a/src/opentslm/time_series_datasets/util.py
+++ b/src/opentslm/time_series_datasets/util.py
@@ -4,23 +4,32 @@
 # SPDX-License-Identifier: MIT
 
 import math
-from typing import List
+from typing import Any, Dict, List
 from opentslm.model_config import PATCH_SIZE
 
 import torch.nn.functional as F
 
 import torch
+from jaxtyping import Float
 
 MAX_VALUE = 50_000
 MIN_VALUE = -MAX_VALUE
 
+# Shape of each element's "time_series" after collation: all of a sample's
+# series stacked and zero-padded to a common length that is a multiple of
+# the patch size.
+PaddedTimeSeries = Float[torch.Tensor, "n_series padded_len"]
+
 
 def extend_time_series_to_match_patch_size_and_aggregate(
-    batch, *, patch_size: int = PATCH_SIZE, normalize: bool = False
-):
+    batch: List[Dict[str, Any]], *, patch_size: int = PATCH_SIZE, normalize: bool = False
+) -> List[Dict[str, Any]]:
     """
     Pad variable-length series so each sample length is a multiple of *patch_size*.
     Optionally normalize each time series to have zero mean and unit variance.
+
+    Each element's ``"time_series"`` entry (a list of 1D array-likes) is replaced
+    by a single ``PaddedTimeSeries`` tensor of shape [n_series, padded_len].
     """
 
     for element in batch:
diff --git a/test/test_jaxtyping_shapes.py b/test/test_jaxtyping_shapes.py
new file mode 100644
index 00000000..6512a8ce
--- /dev/null
+++ b/test/test_jaxtyping_shapes.py
@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
+# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
+#
+# SPDX-License-Identifier: MIT
+
+"""
+Self-contained check that the jaxtyping shape annotations in the codebase are
+actually enforced at runtime.
+
+jaxtyping annotations such as ``Float[torch.Tensor, "*batch patches encoder_dim"]``
+are inert on their own -- Python never validates annotations at call time. They
+only become runtime shape checks when combined with a runtime typechecker
+(beartype here) via jaxtyping's import hook. This test installs that hook for a
+single module and exercises one annotated function with a correct and an
+incorrect shape.
+
+``LinearProjector.forward`` is a good probe because ``torch.nn.Linear`` happily
+accepts a 1-D input (it just treats the last axis as features), so a 1-D tensor
+is rejected *only* by the jaxtyping annotation -- if the bad-shape call raises,
+the annotation is provably being enforced.
+"""
+
+import sys
+import unittest
+
+import torch
+from jaxtyping import TypeCheckError, install_import_hook
+
+TARGET = "opentslm.model.projector.LinearProjector"
+
+# The import hook can only instrument a module imported *after* it is installed,
+# so drop any previously-cached copy to guarantee a fresh, instrumented import.
+for _name in list(sys.modules):
+    if _name == TARGET or _name.startswith(TARGET + "."):
+        del sys.modules[_name]
+
+with install_import_hook(TARGET, "beartype.beartype"):
+    from opentslm.model.projector.LinearProjector import LinearProjector
+
+
+class JaxtypingShapeCheckTest(unittest.TestCase):
+    """Verifies jaxtyping annotations are enforced at runtime via beartype."""
+
+    def setUp(self):
+        self.projector = LinearProjector(input_dim=8, output_dim=16, device="cpu")
+
+    def test_correct_shapes_pass(self):
+        """A well-shaped [*batch, patches, encoder_dim] tensor passes the check."""
+        x = torch.randn(2, 3, 8)  # batch=2, patches=3, encoder_dim=8
+        out = self.projector(x)
+        self.assertEqual(out.shape, (2, 3, 16))
+
+    def test_incorrect_shape_raises(self):
+        """A 1-D tensor lacks the required patches/encoder_dim axes and must be rejected.
+
+        torch.nn.Linear(8, 16) would silently accept this 1-D tensor and return
+        shape [16], so a raised error proves the jaxtyping annotation is enforced
+        rather than torch doing the checking.
+        """
+        bad = torch.randn(8)
+        with self.assertRaises(
+            TypeCheckError,
+            msg="LinearProjector.forward accepted a 1-D tensor; is the beartype "
+            "import hook installed correctly?",
+        ):
+            self.projector(bad)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/uv.lock b/uv.lock
index f1efe58f..16318a5e 100644
--- a/uv.lock
+++ b/uv.lock
@@ -256,6 +256,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b7/b8/3fe70c75fe32afc4bb507f75563d39bc5642255d1d94f1f23604725780bf/babel-2.17.0-py3-none-any.whl", hash = "sha256:4d0b53093fdfb4b21c92b5213dba5a1b23885afa8383709427046b21c366e5f2", size = 10182537, upload-time = "2025-02-01T15:17:37.39Z" },
 ]
 
+[[package]]
+name = "beartype"
+version = "0.22.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c7/94/1009e248bbfbab11397abca7193bea6626806be9a327d399810d523a07cb/beartype-0.22.9.tar.gz", hash = "sha256:8f82b54aa723a2848a56008d18875f91c1db02c32ef6a62319a002e3e25a975f", size = 1608866, upload-time = "2025-12-13T06:50:30.72Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl", hash = "sha256:d16c9bbc61ea14637596c5f6fbff2ee99cbe3573e46a716401734ef50c3060c2", size = 1333658, upload-time = "2025-12-13T06:50:28.266Z" },
+]
+
 [[package]]
 name = "beautifulsoup4"
 version = "4.14.3"
@@ -1039,6 +1048,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7b/55/e5326141505c5d5e34c5e0935d2908a74e4561eca44108fbfb9c13d2911a/isoduration-20.11.0-py3-none-any.whl", hash = "sha256:b2904c2a4228c3d44f409c8ae8e2370eb21a26f7ac2ec5446df141dde3452042", size = 11321, upload-time = "2020-11-01T10:59:58.02Z" },
 ]
 
+[[package]]
+name = "jaxtyping"
+version = "0.3.11"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "wadler-lindig" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/91/c1/091b8852bd7cbf50bd655543c8506033cf4029300c67f8c176c1286879a9/jaxtyping-0.3.11.tar.gz", hash = "sha256:b09c14acf6686feb9e0df5b0d8c6e7c5b6f8d36bf059ee54cd522a186c2ef050", size = 46489, upload-time = "2026-06-13T18:35:23.167Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8b/38/c66bbdc5047f4776c2bd3e47e5295a350e3fa44d5b8942105e71c2a876a0/jaxtyping-0.3.11-py3-none-any.whl", hash = "sha256:8a4bedc4e3f963fa82df41bd13c7ebc2bad925601eb48614c65798f21329d4e3", size = 56593, upload-time = "2026-06-13T18:35:22.01Z" },
+]
+
 [[package]]
 name = "jedi"
 version = "0.19.2"
@@ -1977,7 +1998,7 @@ name = "nvidia-cudnn-cu12"
 version = "9.10.2.21"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cublas-cu12" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" },
@@ -1988,7 +2009,7 @@ name = "nvidia-cufft-cu12"
 version = "11.3.3.83"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink-cu12" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" },
@@ -2015,9 +2036,9 @@ name = "nvidia-cusolver-cu12"
 version = "11.7.3.90"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" },
-    { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" },
-    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cublas-cu12" },
+    { name = "nvidia-cusparse-cu12" },
+    { name = "nvidia-nvjitlink-cu12" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" },
@@ -2028,7 +2049,7 @@ name = "nvidia-cusparse-cu12"
 version = "12.5.8.93"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-nvjitlink-cu12" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" },
@@ -2147,6 +2168,7 @@ dependencies = [
     { name = "datasets" },
     { name = "einops" },
     { name = "huggingface-hub" },
+    { name = "jaxtyping" },
     { name = "matplotlib" },
     { name = "numpy" },
     { name = "open-flamingo" },
@@ -2162,6 +2184,7 @@ dependencies = [
 
 [package.dev-dependencies]
 dev = [
+    { name = "beartype" },
     { name = "reuse" },
     { name = "ruff" },
 ]
@@ -2179,6 +2202,7 @@ requires-dist = [
     { name = "datasets", specifier = ">=2.0" },
     { name = "einops", specifier = ">=0.6" },
     { name = "huggingface-hub", specifier = ">=0.16" },
+    { name = "jaxtyping", specifier = ">=0.2.20" },
     { name = "matplotlib", specifier = ">=3.3" },
     { name = "numpy", specifier = ">=1.21" },
     { name = "open-flamingo", specifier = ">=0.0.2" },
@@ -2194,6 +2218,7 @@ requires-dist = [
 
 [package.metadata.requires-dev]
 dev = [
+    { name = "beartype", specifier = ">=0.22.9" },
     { name = "reuse", specifier = ">=6.2.0" },
     { name = "ruff", specifier = ">=0.14.0" },
 ]
@@ -3705,6 +3730,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bc/56/190ceb8cb10511b730b564fb1e0293fa468363dbad26145c34928a60cb0c/urllib3-2.6.1-py3-none-any.whl", hash = "sha256:e67d06fe947c36a7ca39f4994b08d73922d40e6cca949907be05efa6fd75110b", size = 131138, upload-time = "2025-12-08T15:25:25.51Z" },
 ]
 
+[[package]]
+name = "wadler-lindig"
+version = "0.1.7"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" },
+]
+
 [[package]]
 name = "wandb"
 version = "0.23.1"