From 2f1bd20285a13d3dfa622524f0b34c1b9654d824 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 21:28:07 +0200
Subject: [PATCH 1/2] fix(training): use singular candidate kwargs in test_step

test_step called predict_with_candidates with candidates_x/candidates_y,
but TabR and ModernNCA both define the method as
(self, *data, candidate_x, candidate_y) with no **kwargs, so
trainer.test() on any candidate model raised TypeError. predict_step
already used the correct names.

Part of #413

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 deeptab/training/lightning_module.py |  2 +-
 tests/test_candidate_models.py       | 65 ++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_candidate_models.py

diff --git a/deeptab/training/lightning_module.py b/deeptab/training/lightning_module.py
index 9619add..549595a 100644
--- a/deeptab/training/lightning_module.py
+++ b/deeptab/training/lightning_module.py
@@ -503,7 +503,7 @@ def test_step(self, batch, batch_idx):  # type: ignore
         data, labels = batch
         if hasattr(self.estimator, "predict_with_candidates") and self.train_features is not None:
             preds = self.estimator.predict_with_candidates(  # type: ignore[reportCallIssue]
-                *data, candidates_x=self.train_features, candidates_y=self.train_targets
+                *data, candidate_x=self.train_features, candidate_y=self.train_targets
             )
         else:
             preds = self(*data)
diff --git a/tests/test_candidate_models.py b/tests/test_candidate_models.py
new file mode 100644
index 0000000..fbf3520
--- /dev/null
+++ b/tests/test_candidate_models.py
@@ -0,0 +1,65 @@
+"""Regression test: candidate-model step methods must use the right kwarg names.
+
+``predict_with_candidates`` is defined as ``(self, *data, candidate_x,
+candidate_y)`` in TabR and ModernNCA, but ``test_step`` called it with
+``candidates_x``/``candidates_y``, so ``trainer.test()`` on those models always
+raised TypeError. ``predict_step`` already used the singular names.
+"""
+
+import torch
+import torch.nn as nn
+
+from deeptab.configs import MLPConfig
+from deeptab.training.lightning_module import TaskModel
+
+
+class _CandidateEstimator(nn.Module):
+    """Minimal stand-in with TabR's ``predict_with_candidates`` signature."""
+
+    uses_candidates = True
+    returns_ensemble = False
+
+    def __init__(self, config=None, feature_information=None, num_classes=1, lss=False, **kwargs):
+        super().__init__()
+        self.linear = nn.Linear(4, num_classes)
+        self.seen = {}
+
+    def forward(self, *data):
+        return self.linear(data[0][0])
+
+    def predict_with_candidates(self, *data, candidate_x, candidate_y):
+        self.seen["candidate_x"] = candidate_x
+        self.seen["candidate_y"] = candidate_y
+        return self.forward(*data)
+
+
+def _task_model():
+    task = TaskModel(
+        model_class=_CandidateEstimator,
+        config=MLPConfig(),
+        feature_information=({}, {}, {}),
+        num_classes=1,
+    )
+    task.train_features = ([torch.randn(8, 4)], [], None)
+    task.train_targets = torch.randn(8, 1)
+    return task
+
+
+def _batch(batch_size=4):
+    return (([torch.randn(batch_size, 4)], [], None), torch.randn(batch_size, 1))
+
+
+def test_test_step_passes_candidate_kwargs():
+    task = _task_model()
+    loss = task.test_step(_batch(), 0)
+    assert loss.isfinite()
+    assert task.estimator.seen["candidate_x"] is task.train_features
+    assert task.estimator.seen["candidate_y"] is task.train_targets
+
+
+def test_predict_step_passes_candidate_kwargs():
+    task = _task_model()
+    # predict_step unpacks the batch as the data tuple itself (no labels).
+    preds = task.predict_step(_batch()[0], 0)
+    assert preds.shape == (4, 1)
+    assert task.estimator.seen["candidate_x"] is task.train_features

From 8116d9f38585b03004ea5dff15b8489778470688 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 21:57:43 +0200
Subject: [PATCH 2/2] test: satisfy pyright on the candidate-kwargs regression
 test

compute_loss is typed as returning float, and estimator is typed
nn.Module; cast for the stub's recorded kwargs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 tests/test_candidate_models.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/test_candidate_models.py b/tests/test_candidate_models.py
index fbf3520..6dd316f 100644
--- a/tests/test_candidate_models.py
+++ b/tests/test_candidate_models.py
@@ -6,6 +6,8 @@
 raised TypeError. ``predict_step`` already used the singular names.
 """
 
+from typing import Any, cast
+
 import torch
 import torch.nn as nn
 
@@ -52,9 +54,10 @@ def _batch(batch_size=4):
 def test_test_step_passes_candidate_kwargs():
     task = _task_model()
     loss = task.test_step(_batch(), 0)
-    assert loss.isfinite()
-    assert task.estimator.seen["candidate_x"] is task.train_features
-    assert task.estimator.seen["candidate_y"] is task.train_targets
+    assert torch.as_tensor(loss).isfinite()
+    seen = cast(Any, task.estimator).seen
+    assert seen["candidate_x"] is task.train_features
+    assert seen["candidate_y"] is task.train_targets
 
 
 def test_predict_step_passes_candidate_kwargs():
@@ -62,4 +65,4 @@ def test_predict_step_passes_candidate_kwargs():
     # predict_step unpacks the batch as the data tuple itself (no labels).
     preds = task.predict_step(_batch()[0], 0)
     assert preds.shape == (4, 1)
-    assert task.estimator.seen["candidate_x"] is task.train_features
+    assert cast(Any, task.estimator).seen["candidate_x"] is task.train_features