Fix ps_decoder2 input width when pert_dim is None; add regression test and CI

.github/workflows/ci.yml

@@ -0,0 +1,36 @@
+name: CI
+
+on:
+  push:
+    branches: ["**"]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    name: ubuntu / Python ${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+
+      - name: Install CPU-only torch
+        run: |
+          python -m pip install --upgrade pip
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+      - name: Install pertTF with test extras
+        run: pip install .[test]
+
+      - name: Run tests
+        run: pytest -v

perttf/model/pertTF.py

@@ -119,7 +119,7 @@ def __init__(self,
                 d_model,
                 1,
                 nlayers = ps_decoder2_nlayer,
-                geneinput = self.pert_dim if self.pert_dim is not None else 0,
+                geneinput = pert_dim,  # must match pert_encoder output width (d_model when pert_dim is None)
             )
         else:
             self.ps_decoder2 = None

pyproject.toml

@@ -33,6 +33,9 @@ dependencies = [
     "scanpy >= 1.10.0",
 ]
 
+[project.optional-dependencies]
+test = ["pytest", "torch>=2.2", "tqdm"]
+
 [project.urls]
 Homepage = "https://github.com/davidliwei/pertTF"
 Repository = "https://github.com/davidliwei/pertTF"

tests/test_pertTF_model.py

@@ -0,0 +1,53 @@
+import torch
+
+from perttf.model.pertTF import PerturbationTFModel
+
+D_MODEL, N_PERT, N_PS, NTOKENS, N_CLS = 32, 10, 3, 100, 5
+
+
+def build_model(**overrides):
+    """Mirror the constructor call in notebook/train_pertTF_with__lochNESS.ipynb."""
+    vocab = {"<pad>": 0, **{f"g{i}": i for i in range(1, NTOKENS)}}
+    kwargs = dict(
+        vocab=vocab, dropout=0.0, pad_token="<pad>", pad_value=0,
+        do_mvc=True, do_dab=False, use_batch_labels=False, num_batch_labels=1,
+        domain_spec_batchnorm=False, n_input_bins=0, ecs_threshold=0.7,
+        explicit_zero_prob=False, use_fast_transformer=False, pre_norm=False,
+        n_cls=N_CLS, nlayers_cls=3,
+        pred_lochness_next=True, ps_decoder2_nlayer=5,
+    )
+    kwargs.update(overrides)
+    return PerturbationTFModel(
+        N_PERT, 3, N_PS, NTOKENS, D_MODEL, 4, D_MODEL, 2, **kwargs
+    ).eval()
+
+
+def test_ps_decoder2_input_width_matches_pert_encoder_default():
+    # Regression for #50: with pert_dim unset, the encoder emits d_model-wide
+    # perturbation embeddings, so the decoder must accept 2 * d_model.
+    model = build_model()
+    first_linear = model.ps_decoder2._decoder[0]
+    assert first_linear.in_features == 2 * D_MODEL
+
+
+def test_forward_with_pred_lochness_next_and_default_pert_dim():
+    # The call that crashed in the notebook with
+    # "mat1 and mat2 shapes cannot be multiplied (128x64 and 32x32)".
+    torch.manual_seed(0)
+    model = build_model()
+    B, L = 8, 20
+    src = torch.randint(1, NTOKENS, (B, L))
+    values = torch.rand(B, L)
+    mask = torch.zeros(B, L, dtype=torch.bool)
+    pert = torch.randint(0, N_PERT, (B,))
+    pert_next = torch.randint(0, N_PERT, (B,))
+
+    with torch.no_grad():
+        out = model(
+            src, values, mask,
+            pert_labels=pert, pert_labels_next=pert_next,
+            CLS=True, MVC=True, PERTPRED=True, PSPRED=True,
+        )
+
+    assert out["ps_output_next"].shape == (B, 1)
+    assert torch.isfinite(out["ps_output_next"]).all()