tests: introduce test suite & ci

Author: fel-thomas

.github/workflows/tests.yml

@@ -0,0 +1,20 @@
+name: tests
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install tox torch numpy
+      - name: Run tests
+        run: tox -q

setup.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+
+from setuptools import find_packages, setup
+
+
+def read_version() -> str:
+    version_file = Path(__file__).parent / "raptor" / "__init__.py"
+    for line in version_file.read_text().splitlines():
+        if line.startswith("__version__"):
+            return line.split("=", 1)[1].strip().strip('"').strip("'")
+    raise RuntimeError("Unable to find __version__ in raptor/__init__.py")
+
+
+setup(
+    name="raptor",
+    version=read_version(),
+    description="Block-Recurrent Dynamics in ViTs (Raptor)",
+    packages=find_packages(),
+    python_requires=">=3.9",
+)

tests/conftest.py

@@ -0,0 +1,7 @@
+from pathlib import Path
+import sys
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "raptor"))
+sys.path.insert(0, str(ROOT))

tests/test_block.py

@@ -0,0 +1,22 @@
+import torch
+
+from block import Block
+
+
+def test_block_forward_with_t_scale():
+    block = Block(dim=16, num_heads=4, t_scale=True, swiglu=False)
+    x = torch.randn(2, 6, 16)
+    t = torch.tensor([1.0, 2.0])
+
+    out = block(x, t, t_integer=1)
+
+    assert out.shape == x.shape
+
+
+def test_block_forward_without_t_scale():
+    block = Block(dim=16, num_heads=4, t_scale=False, swiglu=True)
+    x = torch.randn(2, 6, 16)
+
+    out = block(x)
+
+    assert out.shape == x.shape

tests/test_dataloader.py

@@ -0,0 +1,43 @@
+import numpy as np
+import pytest
+import torch
+
+zarr = pytest.importorskip("zarr")
+pytest.importorskip("torchvision")
+pytest.importorskip("PIL")
+
+from dataloader import AsynchZarrLoader, imagenet_transform
+
+
+def test_asynch_zarr_loader_yields_batches(tmp_path):
+    data_path = tmp_path / "data.zarr"
+    data = zarr.open(str(data_path), mode="w", shape=(10, 3, 2, 4), dtype="f4")
+    data[:] = np.random.rand(10, 3, 2, 4).astype("f4")
+
+    loader = AsynchZarrLoader(
+        zarr_path=str(data_path),
+        layer_start=0,
+        layer_end=1,
+        batch_size=4,
+        num_workers=1,
+        queue_size=2,
+        device="cpu",
+    )
+
+    iterator = iter(loader)
+    batch = next(iterator)
+    loader.close()
+    iterator.close()
+
+    assert batch.shape == (4, 2, 2, 4)
+    assert batch.dtype == torch.float32
+
+
+def test_imagenet_transform_output_shape():
+    from PIL import Image
+
+    img = Image.fromarray((np.random.rand(64, 64, 3) * 255).astype("uint8"))
+    tensor = imagenet_transform()(img)
+
+    assert tensor.shape == (3, 224, 224)
+    assert tensor.dtype == torch.float32

tests/test_raptor.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+
+from raptor import Raptor
+
+
+class DummyBlock(nn.Module):
+    def __init__(self, value: float):
+        super().__init__()
+        self.value = value
+
+    def forward(self, x, t):
+        return x + self.value
+
+
+def test_raptor_selects_expected_block():
+    blocks = [DummyBlock(1.0), DummyBlock(2.0), DummyBlock(3.0)]
+    model = Raptor(blocks, thresholds=[2, 5])
+    x = torch.zeros(1, 1, 1)
+    t = torch.tensor([1.0])
+
+    out1 = model(x, t, t_integer=1)
+    out2 = model(x, t, t_integer=3)
+    out3 = model(x, t, t_integer=9)
+
+    assert out1.item() == 1.0
+    assert out2.item() == 2.0
+    assert out3.item() == 3.0

tests/test_scheduler.py

@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from scheduler import CosineScheduler
+
+
+def test_cosine_scheduler_steps_match_schedule():
+    model = torch.nn.Linear(2, 2)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    scheduler = CosineScheduler(
+        optimizer,
+        base_value=0.1,
+        final_value=0.01,
+        total_iters=10,
+        warmup_iters=2,
+    )
+
+    assert len(scheduler.schedule) == 10
+    for _ in range(3):
+        scheduler.step()
+        assert optimizer.param_groups[0]["lr"] == pytest.approx(scheduler[scheduler.iter])
+
+    assert scheduler[10] == pytest.approx(0.01)

tests/test_utils.py

@@ -0,0 +1,27 @@
+import pytest
+
+import utils
+
+
+def test_find_raptor_checkpoint_single_match(tmp_path):
+    models_dir = tmp_path
+    expected = models_dir / "final_raptor3_any_seed_42_step_123.pt"
+    expected.write_text("ok")
+
+    result = utils.find_raptor_checkpoint("raptor3", 42, str(models_dir))
+
+    assert result == str(expected)
+
+
+def test_find_raptor_checkpoint_multiple_matches(tmp_path):
+    models_dir = tmp_path
+    (models_dir / "final_raptor2_a_seed_7_step_1.pt").write_text("ok")
+    (models_dir / "final_raptor2_b_seed_7_step_2.pt").write_text("ok")
+
+    with pytest.raises(ValueError):
+        utils.find_raptor_checkpoint("raptor2", 7, str(models_dir))
+
+
+def test_find_raptor_checkpoint_missing(tmp_path):
+    with pytest.raises(FileNotFoundError):
+        utils.find_raptor_checkpoint("raptor4", 999, str(tmp_path))

tox.ini

@@ -0,0 +1,10 @@
+[tox]
+envlist = py
+skipsdist = true
+
+[testenv]
+deps =
+    pytest
+commands =
+    pytest -q
+sitepackages = true