[ChrisW09]

Fixes #527

Problem

For classification, ContrastivePretrainer.get_knn() keeps only the valid rows: rows that have at least one same-class and one different-class neighbor. contrastive_loss() still repeated every batch row as an anchor. A batch containing a class singleton therefore had more anchor rows than pair rows, and CosineEmbeddingLoss failed with a shape mismatch (32 vs 30). Even with matching sizes, the neighbor rows would have been attributed to the wrong anchors. #414 mentions this crash; its fix corrected the anchor layout but not the anchor selection.

Fix

In deeptab/training/pretraining.py:

• New _pair_indices(labels) returns (knn_indices, neg_indices, anchor_indices), the batch rows each neighbor row belongs to. For regression every row is an anchor.
• get_knn() keeps its public two-value return and delegates to it.
• contrastive_loss(..., anchor_indices=None) builds the anchor block from exactly those rows, in both the pooled and the unpooled branch. A singleton row is not used as an anchor, but it still serves as a negative neighbor for the other rows.
• training_step, validation_step and test_step pass the anchor indices through.

The ValueError for batches in which no row qualifies is unchanged.

Tests

In tests/training/test_training_pretraining.py:

• test_class_singleton_is_dropped_from_anchors_not_from_pairs: anchor selection, the singleton row drawn as a negative for every anchor (with k_neighbors=4 all four different-class rows are drawn, so the check is deterministic), plus finite pooled and unpooled losses.
• test_contrastive_loss_uses_the_given_anchor_rows: with orthonormal embeddings the loss is zero only when the neighbor rows are paired with the given anchors.
• test_classifier_pretraining_completes_with_a_class_singleton: end-to-end FTTransformerClassifier.pretrain() with a class that occurs once.

All three fail on main and pass with this change.

Docs

The pretraining note in docs/api/training/index.rst said a ValueError is raised "if any sample has no same-class or no different-class neighbor". It now says that such samples are skipped as anchors and that the error only occurs when no sample in a batch qualifies. get_knn() and contrastive_loss() have numpydoc Parameters/Returns sections; get_knn() explains that its rows need not match the batch rows instead of pointing to the private _pair_indices().

Checks

• just test: 3239 passed, 46 skipped, 28 xfailed; coverage 91.56%
• just docs: build succeeded (warnings treated as errors)
• just check: all hooks passed (ruff, prettier, pyright)
• The single-row-batch crash is a separate bug ([BUG] Contrastive pretraining crashes on a batch too small to form pairs (e.g. a one-row last batch) #529), fixed in fix(training): skip pretraining batches too small to form pairs #530. I checked that the two branches merge cleanly together.