Add jaxtyping, fix minor bugs thus uncovered by TimoStoff · Pull Request #50 · SchmiedmayerLab/OpenTSLM

TimoStoff · 2026-07-07T10:59:58Z

Add jaxtyping runtime shape-checking

♻️ Current situation & Problem

Tensor shape mismatches in the model/data pipeline (e.g. batch dimensions, media/time-series axes) were previously only caught implicitly at runtime via shape-mismatch errors deep in torch or open_flamingo internals, with no static or explicit documentation of expected tensor shapes. This made bugs like the Flamingo cross-attention media-axis misalignment (see below) hard to spot from reading the code alone.

⚙️ Release Notes

Added jaxtyping and beartype as dependencies for runtime tensor shape/type checking.
Introduced an OPENTSLM_MODE=dev environment gate that enables jaxtyping/beartype runtime checks in development without imposing their overhead in production.
Annotated tensor-producing/consuming functions across the model stack with jaxtyping shape annotations. Basically any file with functions taking or producing tensor data.
Added a ruff ignore for rule F722 (syntax-error false positives on jaxtyping's string-based shape annotations).
Bug fix: Fixed a Flamingo cross-attention bug where OpenTSLMFlamingo packed all of a sample's time series into a single media chunk (unsqueeze(1) → [B, 1, n_series, L], i.e. T_media=1) even though the prompt emits one <image> token per series. Since Flamingo's masked cross-attention aligns each text token to the media chunk matching its running <image> count, this silently collapsed cross-attention to a uniform average over latents for every multi-series sample (e.g. ECG-QA's 12 leads, HAR's 3 axes). Fixed by using unsqueeze(2) → [B, n_series, 1, L] so each series is its own media chunk (T_media=n_series).
Fixed a constructor argument mismatch in TransformerMLPEncoder.
Fixed several any → Any typos in batch type hints (the built-in any was being used where the typing.Any type was intended).

📚 Documentation

Shape and type expectations are now documented inline via jaxtyping annotations directly on function signatures, rather than in prose . e.g. Float[torch.Tensor, "batch n_series length"]. This keeps the shape contract next to the code it constrains and lets it be checked automatically rather than only described in comments.

✅ Testing

Added a unit test (test_jaxtyping_shapes) exercising a simple jaxtyping runtime shape check to confirm the OPENTSLM_MODE=dev gate correctly enables enforcement.
Added a small smoke-test training run to confirm the annotated pipeline still trains end-to-end after the shape-annotation and Flamingo fix changes.

Code of Conduct & Contributing Guidelines

By creating and submitting this pull request, you agree to follow our Code of Conduct and Contributing Guidelines:

I agree to follow the Code of Conduct and Contributing Guidelines.

^{Need help on this PR? Tag /codesmith with what you need. Autofix is disabled.}

Summary by CodeRabbit

New Features
- Added a lightweight smoke test mode for curriculum training to run a single epoch and a minimal number of batches, making quick validation faster.
- Introduced a development mode that enables stricter runtime shape/type checks for model components.
Tests
- Added coverage to verify tensor shape checks are enforced at runtime.
Chores
- Updated dependencies and linting settings to support the new typing and validation improvements.

Description Jaxtyping is a really cool library that allows adding type annotations to python functions that accept tensor inputs. I really like using this because it helps to find otherwise subtle and hard to detect bugs. Adding it to the deps will allow adding such typing annotations in the following commits. More details here: https://github.com/patrick-kidger/jaxtyping Test Plan None for now, but will be added in the following commits.

Description Beartype will be needed to enable jaxtyping to be more than just code aesthetic. It will allow me to enable runtime type checking for jaxtyping annos. Test Plan Empty again, but will be tested in downstream commits. Just keeping atomic changes split for now.

Description Install jaxtyping's beartype import hook for the whole opentslm package when OPENTSLM_MODE=dev, so Float[...] shape annotations and ordinary type hints are enforced at runtime across every module. Defaults to "opt": annotations stay inert (zero overhead) and beartype is not required to import the package. Test Plan Noop commit. I ran ``` CUDA_VISIBLE_DEVICES="" OPENTSLM_MODE=dev\ uv run python curriculum_learning.py \ --model OpenTSLMSP \ --llm_id google/gemma-3-270m \ --stages stage1_mcq \ --device cpu \ --batch_size 2 ``` to double check.

Description: jaxtyping shape strings like Float[Tensor, "batch length"] are misread by pyflakes as malformed forward references, producing 30 spurious F722 errors across the annotated files. Add the standard jaxtyping ruff ignore. This is a pure Claude fix. Test Plan None

Description Adds jaxtyping annotations. As in the title. Test Plan Ran training at the top of this stack with ``` OPENTSLM_MODE=dev uv run python curriculum_learning.py \ --model OpenTSLMSP \ --llm_id google/gemma-3-270m \ --stages stage1_mcq \ --device cpu \ --batch_size 2 \ --smoke_test ``` output: `torch_dtype` is deprecated! Use `dtype` instead! The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False` 🎓 Starting Curriculum Learning with OpenTSLMSP 📊 All stages: stage1_mcq 🔄 Incomplete stages: stage1_mcq 💻 Device: cpu 📦 Batch size: 2 ================================================================================ 🚀 Starting stage1_mcq Training with OpenTSLMSP ============================================================ 📊 Stage Configuration: Epochs: 1 Encoder LR: 2.00e-04 Projector LR: 1.00e-04 Batch size per GPU: 2 🆕 Starting with fresh model (first stage) ℹ️ LoRA disabled for stage1_mcq (only enabled for stages 3+) 📊 Learning rates for OpenTSLMSP: Encoder LR: 2.00e-04 Projector LR: 1.00e-04 Formatting training samples... Training samples: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 38400/38400 [00:18<00:00, 2124.24it/s] Formatting validation samples... Validation samples: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4800/4800 [00:02<00:00, 2171.36it/s] Formatting test samples... Test samples: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4800/4800 [00:02<00:00, 2111.12it/s] 📈 Total training steps: 19200 🔥 Warmup steps: 576 🆕 Starting fresh training for stage1_mcq Epoch 1/1: 0%| | 0/19200 [00:00<?, ?it/s][DEBUG] Batch 0 - batch size: 2 [DEBUG] Sample key 'time_series' shape: torch.Size([1, 128]) [DEBUG] Sample key 'time_series_text' list length: 1 No CUDA Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation. Epoch 1/1: 0%| | 0/19200 [14:13<?, ?it/s, loss=4.4998, lr=3.47e-07] Epoch 1 — train loss: 4.4998 Validating stage1_mcq: 0%| | 0/4800 [00:02<?, ?it/s] Epoch 1 — val loss: 3.9003 Epoch 1 — best loss: inf 💾 Disk space: 92.80 GB free in results/gemma_3_270m/OpenTSLMSP/stage1_mcq/checkpoints ✔️ New best model saved. 📂 Loaded best checkpoint from epoch 1 for evaluation. 🏁 Training completed for stage1_mcq Total epochs run: 1 Best validation loss: 3.9003 Epochs without improvement: 0 [Eval] rank=0, world_size=1 Saving per-rank test predictions to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions_rank_0.jsonl Evaluating stage1_mcq: 0%| | 0/4800 [03:01<?, ?it/s] Merged per-rank predictions into: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions.jsonl ✅ stage1_mcq evaluation complete: Test predictions saved to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions.jsonl Metrics saved to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/metrics.json Max tokens used for generation: 2000 test_loss: nan epoch: 1.0000 accuracy: 0.0000 ✅ Stage stage1_mcq marked as completed 🎉 Curriculum Learning Complete! 📁 All results saved to: results/gemma_3_270m/OpenTSLMSP/ 📊 Overall results: results/gemma_3_270m/OpenTSLMSP/curriculum_results.json ```

Adds jaxtyping annotations. As in the title. Test Plan Ran training at the top of this stack with ``` OPENTSLM_MODE=dev uv run python curriculum_learning.py \ --model OpenTSLMSP \ --llm_id google/gemma-3-270m \ --stages stage1_mcq \ --device cpu \ --batch_size 2 \ --smoke_test ``` output: `torch_dtype` is deprecated! Use `dtype` instead! The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False` 🎓 Starting Curriculum Learning with OpenTSLMSP 📊 All stages: stage1_mcq 🔄 Incomplete stages: stage1_mcq 💻 Device: cpu 📦 Batch size: 2 ================================================================================ 🚀 Starting stage1_mcq Training with OpenTSLMSP ============================================================ 📊 Stage Configuration: Epochs: 1 Encoder LR: 2.00e-04 Projector LR: 1.00e-04 Batch size per GPU: 2 🆕 Starting with fresh model (first stage) ℹ️ LoRA disabled for stage1_mcq (only enabled for stages 3+) 📊 Learning rates for OpenTSLMSP: Encoder LR: 2.00e-04 Projector LR: 1.00e-04 Formatting training samples... Training samples: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 38400/38400 [00:18<00:00, 2124.24it/s] Formatting validation samples... Validation samples: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4800/4800 [00:02<00:00, 2171.36it/s] Formatting test samples... Test samples: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4800/4800 [00:02<00:00, 2111.12it/s] 📈 Total training steps: 19200 🔥 Warmup steps: 576 🆕 Starting fresh training for stage1_mcq Epoch 1/1: 0%| | 0/19200 [00:00<?, ?it/s][DEBUG] Batch 0 - batch size: 2 [DEBUG] Sample key 'time_series' shape: torch.Size([1, 128]) [DEBUG] Sample key 'time_series_text' list length: 1 No CUDA Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation. Epoch 1/1: 0%| | 0/19200 [14:13<?, ?it/s, loss=4.4998, lr=3.47e-07] Epoch 1 — train loss: 4.4998 Validating stage1_mcq: 0%| | 0/4800 [00:02<?, ?it/s] Epoch 1 — val loss: 3.9003 Epoch 1 — best loss: inf 💾 Disk space: 92.80 GB free in results/gemma_3_270m/OpenTSLMSP/stage1_mcq/checkpoints ✔️ New best model saved. 📂 Loaded best checkpoint from epoch 1 for evaluation. 🏁 Training completed for stage1_mcq Total epochs run: 1 Best validation loss: 3.9003 Epochs without improvement: 0 [Eval] rank=0, world_size=1 Saving per-rank test predictions to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions_rank_0.jsonl Evaluating stage1_mcq: 0%| | 0/4800 [03:01<?, ?it/s] Merged per-rank predictions into: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions.jsonl ✅ stage1_mcq evaluation complete: Test predictions saved to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions.jsonl Metrics saved to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/metrics.json Max tokens used for generation: 2000 test_loss: nan epoch: 1.0000 accuracy: 0.0000 ✅ Stage stage1_mcq marked as completed 🎉 Curriculum Learning Complete! 📁 All results saved to: results/gemma_3_270m/OpenTSLMSP/ 📊 Overall results: results/gemma_3_270m/OpenTSLMSP/curriculum_results.json ```

Description Adds jaxtyping annotations. As in the title. Test Plan Ran training at the top of this stack with ``` OPENTSLM_MODE=dev uv run python curriculum_learning.py \ --model OpenTSLMSP \ --llm_id google/gemma-3-270m \ --stages stage1_mcq \ --device cpu \ --batch_size 2 \ --smoke_test ``` output: `torch_dtype` is deprecated! Use `dtype` instead! The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False` 🎓 Starting Curriculum Learning with OpenTSLMSP 📊 All stages: stage1_mcq 🔄 Incomplete stages: stage1_mcq 💻 Device: cpu 📦 Batch size: 2 ================================================================================ 🚀 Starting stage1_mcq Training with OpenTSLMSP ============================================================ 📊 Stage Configuration: Epochs: 1 Encoder LR: 2.00e-04 Projector LR: 1.00e-04 Batch size per GPU: 2 🆕 Starting with fresh model (first stage) ℹ️ LoRA disabled for stage1_mcq (only enabled for stages 3+) 📊 Learning rates for OpenTSLMSP: Encoder LR: 2.00e-04 Projector LR: 1.00e-04 Formatting training samples... Training samples: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 38400/38400 [00:18<00:00, 2124.24it/s] Formatting validation samples... Validation samples: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4800/4800 [00:02<00:00, 2171.36it/s] Formatting test samples... Test samples: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4800/4800 [00:02<00:00, 2111.12it/s] 📈 Total training steps: 19200 🔥 Warmup steps: 576 🆕 Starting fresh training for stage1_mcq Epoch 1/1: 0%| | 0/19200 [00:00<?, ?it/s][DEBUG] Batch 0 - batch size: 2 [DEBUG] Sample key 'time_series' shape: torch.Size([1, 128]) [DEBUG] Sample key 'time_series_text' list length: 1 No CUDA Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation. Epoch 1/1: 0%| | 0/19200 [14:13<?, ?it/s, loss=4.4998, lr=3.47e-07] Epoch 1 — train loss: 4.4998 Validating stage1_mcq: 0%| | 0/4800 [00:02<?, ?it/s] Epoch 1 — val loss: 3.9003 Epoch 1 — best loss: inf 💾 Disk space: 92.80 GB free in results/gemma_3_270m/OpenTSLMSP/stage1_mcq/checkpoints ✔️ New best model saved. 📂 Loaded best checkpoint from epoch 1 for evaluation. 🏁 Training completed for stage1_mcq Total epochs run: 1 Best validation loss: 3.9003 Epochs without improvement: 0 [Eval] rank=0, world_size=1 Saving per-rank test predictions to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions_rank_0.jsonl Evaluating stage1_mcq: 0%| | 0/4800 [03:01<?, ?it/s] Merged per-rank predictions into: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions.jsonl ✅ stage1_mcq evaluation complete: Test predictions saved to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/test_predictions.jsonl Metrics saved to: results/gemma_3_270m/OpenTSLMSP/stage1_mcq/results/metrics.json Max tokens used for generation: 2000 test_loss: nan epoch: 1.0000 accuracy: 0.0000 ✅ Stage stage1_mcq marked as completed 🎉 Curriculum Learning Complete! 📁 All results saved to: results/gemma_3_270m/OpenTSLMSP/ 📊 Overall results: results/gemma_3_270m/OpenTSLMSP/curriculum_results.json ```

Description Simple unit test just to check that jaxtyping is working as intended. Test Plan ``` uv run python -m unittest test.test_jaxtyping_shapes -v ```

The batch parameters were annotated List[Dict[str, any]] using the builtin any (a function) instead of typing.Any. This is invalid as a type hint and made the jaxtyping import hook (jaxtyping.install_import_hook + beartype) fail at import time, blocking package-wide runtime shape checking. Import typing.Any and use it. Replicate the check (requires jaxtyping + beartype installed; run from repo root). Before this fix it raised BeartypeDecorHintNonpepException; now it prints the OK line: PYTHONPATH=src python -c " from jaxtyping import install_import_hook with install_import_hook('opentslm', 'beartype.beartype'): import opentslm # transitively imports OpenTSLMSP and OpenTSLMFlamingo print('OK: opentslm decorated cleanly by beartype (no invalid type hints)') "

super().__init__ passed 3 positional args (input_dim, output_dim, dropout) but TimeSeriesEncoderBase.__init__ accepts only (output_dim, dropout), and self.input_dim was never set despite being used in forward. The encoder could not be instantiated at all; it is currently unused, so this went unnoticed. Mirror the TransformerCNNEncoder pattern: pass (output_dim, dropout) to super and set self.input_dim in the subclass. Surfaced by running under OPENTSLM_MODE=dev.

I added the ability to run small test runs so that I could actually test end-to-end changes on my puny local laptop. Run it with ``` CUDA_VISIBLE_DEVICES="" OPENTSLM_MODE=dev uv run python curriculum_learning.py \ --model OpenTSLMSP \ --llm_id google/gemma-3-270m \ --stages stage1_mcq \ --device cpu \ --batch_size 2 \ --smoke_test ```

coderabbitai · 2026-07-07T11:00:19Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 208f631c-cf64-4cb2-a946-4e8bde19e518

📥 Commits

Reviewing files that changed from the base of the PR and between 104013b and 1f6d481.

⛔ Files ignored due to path filters (1)

uv.lock is excluded by !**/*.lock

📒 Files selected for processing (19)

curriculum_learning.py
pyproject.toml
requirements.txt
src/opentslm/__init__.py
src/opentslm/model/encoder/CNNTokenizer.py
src/opentslm/model/encoder/TimeSeriesEncoderBase.py
src/opentslm/model/encoder/TransformerCNNEncoder.py
src/opentslm/model/encoder/TransformerMLPEncoder.py
src/opentslm/model/llm/OpenTSLMFlamingo.py
src/opentslm/model/llm/OpenTSLMSP.py
src/opentslm/model/llm/TimeSeriesFlamingoWithTrainableEncoder.py
src/opentslm/model/llm/TimeSeriesLLM.py
src/opentslm/model/projector/LinearProjector.py
src/opentslm/model/projector/MLPProjector.py
src/opentslm/prompt/text_time_series_prompt.py
src/opentslm/time_series_datasets/monash/monash_utils.py
src/opentslm/time_series_datasets/ucr/ucr_loader.py
src/opentslm/time_series_datasets/util.py
test/test_jaxtyping_shapes.py

📝 Walkthrough

Walkthrough

This PR adds jaxtyping-based runtime tensor shape annotations across encoders, projectors, LLM wrappers, and dataset utilities, enabled via a beartype import hook gated by an environment variable. It also adds a smoke_test mode to the curriculum training script for single-batch/epoch fast runs, plus new dependencies and a shape-enforcement test.

Changes

Jaxtyping Runtime Shape Checking

Layer / File(s)	Summary
Dependencies and lint config `pyproject.toml`, `requirements.txt`	Adds `jaxtyping` and `beartype` dependencies and configures ruff to ignore F722 for jaxtyping string annotations.
Import hook installation `src/opentslm/__init__.py`	Conditionally installs a beartype/jaxtyping import hook when `OPENTSLM_MODE=dev`.
Encoder shape annotations `src/opentslm/model/encoder/CNNTokenizer.py`, `TimeSeriesEncoderBase.py`, `TransformerCNNEncoder.py`, `TransformerMLPEncoder.py`	Forward methods gain `Float[...]` shape annotations; `TransformerMLPEncoder` adjusts its base-class init call and stores `input_dim`.
Projector shape annotations and test `src/opentslm/model/projector/LinearProjector.py`, `MLPProjector.py`, `test/test_jaxtyping_shapes.py`	Projector forward methods add `Float[...]` annotations; a new test verifies shape enforcement raises `TypeCheckError` for incorrect shapes.
LLM wrapper typing and reshape `src/opentslm/model/llm/OpenTSLMFlamingo.py`, `OpenTSLMSP.py`, `TimeSeriesFlamingoWithTrainableEncoder.py`, `TimeSeriesLLM.py`	Batch/return types updated to `Any`/`jaxtyping`; `OpenTSLMFlamingo.pad_and_apply_batch` changes media tensor reshape from `unsqueeze(1)` to `unsqueeze(2)`.
Dataset and prompt shape annotations `src/opentslm/prompt/text_time_series_prompt.py`, `src/opentslm/time_series_datasets/monash/monash_utils.py`, `ucr/ucr_loader.py`, `util.py`	Adds `Num`/`Shaped`/`Float`/`Int` shape annotations and a new `PaddedTimeSeries` type alias.

Curriculum Training Smoke Test Mode

Layer / File(s)	Summary
Smoke-test flag and wiring `curriculum_learning.py`	Adds `--smoke_test` CLI flag and a `smoke_test` constructor parameter passed into `CurriculumTrainer`.
Smoke-test loop short-circuiting `curriculum_learning.py`	Training/evaluation loops cap to one epoch/batch when `smoke_test` is enabled; loss averaging uses tracked batch counters.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant CLI
  participant CurriculumTrainer
  participant TrainLoop as _train_stage
  participant EvalLoop as _evaluate_stage

  CLI->>CurriculumTrainer: smoke_test=True
  CurriculumTrainer->>TrainLoop: num_epochs=1
  TrainLoop->>TrainLoop: break after 1 batch, track num_train_batches
  TrainLoop->>EvalLoop: run validation
  EvalLoop->>EvalLoop: increment num_val_batches, break after 1 batch
  EvalLoop-->>TrainLoop: avg_val_loss

sequenceDiagram
  participant Test as JaxtypingShapeCheckTest
  participant Hook as BeartypeImportHook
  participant Projector as LinearProjector

  Test->>Hook: install_import_hook(LinearProjector module)
  Test->>Projector: import and instantiate
  Test->>Projector: forward(correct shape tensor)
  Projector-->>Test: valid output
  Test->>Projector: forward(incorrect shape tensor)
  Projector-->>Test: raises TypeCheckError

Related PRs: None identified from the provided context.

Suggested labels: enhancement, dependencies, testing

Suggested reviewers: Reviewers familiar with the model encoder/projector/LLM modules and the curriculum training pipeline.

🐰 A hop through tensors, shapes now checked with care,
Beartype whispers warnings when dimensions aren't fair,
One smoky batch to test the training grind,
Jaxtyping annotations, precisely aligned.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title matches the PR’s main theme: adding jaxtyping/runtime checks and fixing bugs uncovered during that work.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

TimoStoff added 22 commits July 6, 2026 12:48

Add jaxtyping runtime shape-check unit test

6ecbe7a

Description Simple unit test just to check that jaxtyping is working as intended. Test Plan ``` uv run python -m unittest test.test_jaxtyping_shapes -v ```

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Add jaxtyping, fix minor bugs thus uncovered#50

Add jaxtyping, fix minor bugs thus uncovered#50
TimoStoff wants to merge 22 commits into
SchmiedmayerLab:mainfrom
TimoStoff:add-jaxtyping

TimoStoff commented Jul 7, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Uh oh!

Conversation

TimoStoff commented Jul 7, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Add jaxtyping runtime shape-checking

♻️ Current situation & Problem

⚙️ Release Notes

📚 Documentation

✅ Testing

Code of Conduct & Contributing Guidelines

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

TimoStoff commented Jul 7, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 7, 2026 •

edited

Loading