[BUG] TabR/ModernNCA leak query labels into training-time retrieval; TabR memory_efficient path broken

[ChrisW09]

Describe the bug
Retrieval-based models leak the query's own label into training.

TabR (deeptab/architectures/tabr.py:320-340): the original TabR implementation prepends the batch to the candidate pool before searching, which is what makes the anti-leakage mask context_idx == arange(batch_size) correct. This port skips the prepend — TaskModel.setup/training_step (deeptab/training/lightning_module.py:265-293, 395-401) pass the entire training set, unmodified, as candidate_x/candidate_y. Consequences:

• Query row i exists in the pool at its dataset index j ≠ i, survives the mask, and is retrieved at distance ≈ 0 → highest softmax weight → its own label embedding (context_y_emb) feeds the prediction. Training collapses toward label copying; validation (no self-match) sees a distribution shift.
• Meanwhile the mask removes candidates 0..B-1 — legitimate neighbors — at random.
• The targets argument that exists precisely for the prepend ("Targets for the query rows, concatenated with the candidate pool", tabr.py:274-275) is asserted non-None and then never used.

TabR memory_efficient=True (tabr.py:342-347) is additionally broken: x has already been reassigned to the encoded (B, d_main) tensor at line 314, so torch.cat([x, candidate_x]) concatenates encoded queries with raw (N, input_dim) features (shape error unless dims coincide), and context_idx — indices into the candidate-only pool — is offset by batch_size against the concatenated tensor.

ModernNCA (deeptab/architectures/experimental/modern_nca.py:193-207): correctly prepends the batch and masks the diagonal, but because the trainer passes the full training set as candidates, each query's duplicate row can still be drawn by torch.randperm(data_size)[:retrieval_size] at an off-diagonal position — with the default sample_rate=0.5, roughly half of training queries see their own row at distance 0. The reference implementation excludes the current batch from the pool.

To Reproduce
Read train_with_candidates in tabr.py against the reference TabR (candidate_k = torch.cat([k, candidate_k]); candidate_y = torch.cat([y, candidate_y]) before the faiss search). Empirically: train TabR on data where y is pure noise — train loss drops far below what a no-leak model can achieve while val loss doesn't move.

Expected behavior
Training-time candidate retrieval must exclude the query row: either prepend the batch and keep the existing mask (reference behavior), or exclude batch indices from the pool before the search. The memory_efficient path should re-encode raw candidate rows only, with correctly offset indices.

Screenshots
n/a

Desktop (please complete the following information):

• OS: macOS (Darwin 25.5.0, arm64)
• Python version: 3.11.15
• deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
Related wiring note: TaskModel.test_step (lightning_module.py:504-507) calls predict_with_candidates(candidates_x=..., candidates_y=...) but both implementations take candidate_x/candidate_y (singular) → trainer.test() on TabR/ModernNCA always raises TypeError. predict_step uses the correct names.