fix(core): second-pass review findings

deeptab/core/base_model.py

@@ -267,7 +267,6 @@ def encode(self, data, grad=False):
                 if available_layer == "rnn":
                     embeddings, _ = layer(x)  # type: ignore[reportCallIssue]
                 else:
-                    embeddings = self.encoder(x)  # type: ignore[reportCallIssue]
                     embeddings = layer(x)  # type: ignore[reportCallIssue]
         else:
             x = self.embedding_layer(*data)  # type: ignore[reportCallIssue]

deeptab/models/base.py

@@ -98,7 +98,7 @@ def _validate_fit_inputs(
         family_lower = family.lower()
         if family_lower in {"poisson", "negativebinom"} and (y_arr < 0).any():
             raise target_range_error(family, "non-negative")
-        if family_lower in {"gamma", "inversegaussian"} and (y_arr <= 0).any():
+        if family_lower in {"gamma", "inversegamma", "lognormal"} and (y_arr <= 0).any():
             raise target_range_error(family, "strictly positive")
         if family_lower == "binomial" and not np.all((y_arr == 0) | (y_arr == 1)):
             raise target_range_error(family, "binary (0 or 1)")

deeptab/nn/blocks/common.py

@@ -83,7 +83,10 @@ def forward(ctx, input_, dim=-1):
         """
         ctx.dim = dim
         max_val, _ = input_.max(dim=dim, keepdim=True)
-        input_ -= max_val  # Numerical stability trick, as with softmax.
+        # Numerical stability trick, as with softmax. Must not be in-place:
+        # callers pass Parameters directly (e.g. ODST feature_selection_logits)
+        # and an in-place shift would corrupt the stored values.
+        input_ = input_ - max_val
         tau, supp_size = SparsemaxFunction._threshold_and_support(input_, dim=dim)
         output = torch.clamp(input_ - tau, min=0)
         ctx.save_for_backward(supp_size, output)

deeptab/training/lightning_module.py

@@ -201,6 +201,11 @@ def __init__(
         self.early_pruning_threshold = early_pruning_threshold
         self.pruning_epoch = pruning_epoch
         self.val_losses = []
+        # Candidate pools for retrieval models; populated in setup("fit").
+        # Must default to None so the `is not None` guards in the step
+        # methods work for models restored without an in-process fit.
+        self.train_features = None
+        self.train_targets = None
 
         # Store custom metrics
         self.train_metrics = train_metrics or {}

tests/test_second_pass_fixes.py

@@ -0,0 +1,58 @@
+"""Regression tests for the second-pass review findings."""
+
+from typing import cast
+
+import pandas as pd
+import pytest
+import torch
+
+from deeptab.configs import MLPConfig
+from deeptab.distributions.registry import DISTRIBUTION_REGISTRY
+from deeptab.models.base import _validate_fit_inputs
+from deeptab.nn.blocks.common import sparsemax
+from deeptab.training.lightning_module import TaskModel
+
+
+class _Dummy(torch.nn.Module):
+    def __init__(self, config=None, feature_information=None, num_classes=1, lss=False, **kwargs):
+        super().__init__()
+        self.linear = torch.nn.Linear(2, num_classes)
+
+    def forward(self, *data):
+        return self.linear(data[0])
+
+
+def test_sparsemax_does_not_mutate_input():
+    """ODST passes its feature_selection_logits Parameter straight in."""
+    logits = torch.randn(4, 8)
+    original = logits.clone()
+    sparsemax(logits, dim=-1)
+    assert torch.equal(logits, original)
+
+
+def test_sparsemax_output_unchanged_by_the_fix():
+    logits = torch.tensor([[1.0, 2.0, 3.0], [0.0, 0.0, 5.0]])
+    out = cast(torch.Tensor, sparsemax(logits, dim=-1))
+    assert torch.allclose(out.sum(dim=-1), torch.ones(2), atol=1e-6)
+    assert (out >= 0).all()
+
+
+class TestValidateFitInputsFamilies:
+    @pytest.mark.parametrize("family", ["gamma", "inversegamma", "lognormal"])
+    def test_strictly_positive_families(self, family):
+        X = pd.DataFrame({"a": [1.0, 2.0, 3.0]})
+        y = [1.0, 0.0, 2.0]
+        with pytest.raises(Exception, match="positive"):
+            _validate_fit_inputs(X, y, regression=True, family=family)
+
+    def test_checked_families_exist_in_the_registry(self):
+        """The old check gated on 'inversegaussian', which is not a family."""
+        assert "inversegamma" in DISTRIBUTION_REGISTRY
+        assert "inversegaussian" not in DISTRIBUTION_REGISTRY
+
+
+def test_taskmodel_candidate_pool_defaults_to_none():
+    """The `is not None` guards presume a None default, set only in setup('fit')."""
+    task = TaskModel(model_class=_Dummy, config=MLPConfig(), feature_information=({}, {}, {}))
+    assert task.train_features is None
+    assert task.train_targets is None