From 464d55dd5e3c14f83a909af7bc04fe687fcf00f3 Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:34:51 +0200 Subject: [PATCH 1/2] fix(core): second-pass review findings - BaseModel.encode(): drop the stray self.encoder(x) call in the no-grad branch. Mamba architectures have no 'encoder' attribute, so the public clf.encode(X) path raised AttributeError for them; for every other model it ran the encoder twice and discarded the first result. The grad=True branch never had the extra call. - _validate_fit_inputs gated the strictly-positive check on 'inversegaussian', which is not in DISTRIBUTION_REGISTRY (the real key is 'inversegamma'), so the check never fired; lognormal needs it too. - sparsemax shifted its input in place. ODST passes the feature_selection_logits Parameter directly, so the stored parameter values were corrupted by -max on first use. Output is shift-invariant, so predictions were unaffected, but weight decay and checkpoints saw the shifted values. - TaskModel now initialises train_features/train_targets to None, so the 'is not None' guards in the step methods work for candidate models restored from a checkpoint without an in-process fit. Fixes #426 Co-Authored-By: Claude Fable 5 --- deeptab/core/base_model.py | 1 - deeptab/models/base.py | 2 +- deeptab/nn/blocks/common.py | 5 ++- deeptab/training/lightning_module.py | 5 +++ tests/test_second_pass_fixes.py | 56 ++++++++++++++++++++++++++++ 5 files changed, 66 insertions(+), 3 deletions(-) create mode 100644 tests/test_second_pass_fixes.py diff --git a/deeptab/core/base_model.py b/deeptab/core/base_model.py index d86e85a9..abc787e9 100644 --- a/deeptab/core/base_model.py +++ b/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] diff --git a/deeptab/models/base.py b/deeptab/models/base.py index 0301bbfb..2c6a6594 100644 --- a/deeptab/models/base.py +++ b/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)") diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py index 0ada7006..26675e1d 100644 --- a/deeptab/nn/blocks/common.py +++ b/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) diff --git a/deeptab/training/lightning_module.py b/deeptab/training/lightning_module.py index 9619add7..146e48ec 100644 --- a/deeptab/training/lightning_module.py +++ b/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 {} diff --git a/tests/test_second_pass_fixes.py b/tests/test_second_pass_fixes.py new file mode 100644 index 00000000..d85feb71 --- /dev/null +++ b/tests/test_second_pass_fixes.py @@ -0,0 +1,56 @@ +"""Regression tests for the second-pass review findings.""" + +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 = 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 From b9bb4b2dfee582f0f6d618fbc29961d6a7fbb6a3 Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:01:01 +0200 Subject: [PATCH 2/2] test: satisfy pyright on the second-pass regression tests sparsemax is typed as returning Tensor | None via the autograd Function. Co-Authored-By: Claude Fable 5 --- tests/test_second_pass_fixes.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_second_pass_fixes.py b/tests/test_second_pass_fixes.py index d85feb71..700ac21e 100644 --- a/tests/test_second_pass_fixes.py +++ b/tests/test_second_pass_fixes.py @@ -1,5 +1,7 @@ """Regression tests for the second-pass review findings.""" +from typing import cast + import pandas as pd import pytest import torch @@ -30,7 +32,7 @@ def test_sparsemax_does_not_mutate_input(): def test_sparsemax_output_unchanged_by_the_fix(): logits = torch.tensor([[1.0, 2.0, 3.0], [0.0, 0.0, 5.0]]) - out = sparsemax(logits, dim=-1) + 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()