From 8eb0095b1494dbe818c99dacae3aa2c93a1b323f Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:40:09 +0200 Subject: [PATCH] fix: Multinomial probabilities, sparsemax gradients, Tangos single-row batches Three independent defects: - MultinomialDistribution.forward returned raw logits. BaseDistribution .forward looks up a '_transform' attribute per parameter, but this family names its parameters per class (p_0, p_1, ...) and keeps a single shared 'probs_transform', so every lookup missed and fell back to identity. predict(raw=False) therefore returned logits labelled as probabilities (e.g. [2.0, 0.0, -1.0]) while compute_loss applied the softmax correctly. Overridden to apply the shared transform. - SparsemaxFunction.backward divided by supp_size.squeeze(), and a bare squeeze() also drops unrelated size-1 axes; the resulting shape mismatch broadcast into silently wrong gradients whenever a non-reduced dimension had size 1 (e.g. NODE with depth 1). Now squeeze(dim), which only removes the reduced axis. - Tangos.penalty_forward called jacobian.squeeze(), which removes the batch axis when the batch holds a single row -- routine for the final batch of an odd-sized dataset -- so the following neuron_attr.shape[2] raised IndexError. Singleton axes other than the batch axis are now squeezed individually. Fixes #453 Co-Authored-By: Claude Fable 5 --- deeptab/architectures/experimental/tangos.py | 5 +- deeptab/distributions/categorical.py | 10 +++ deeptab/nn/blocks/common.py | 4 +- tests/test_distribution_misc_fixes.py | 76 ++++++++++++++++++++ 4 files changed, 93 insertions(+), 2 deletions(-) create mode 100644 tests/test_distribution_misc_fixes.py diff --git a/deeptab/architectures/experimental/tangos.py b/deeptab/architectures/experimental/tangos.py index 3f899cb0..8a6ae891 100644 --- a/deeptab/architectures/experimental/tangos.py +++ b/deeptab/architectures/experimental/tangos.py @@ -187,7 +187,10 @@ def penalty_forward(self, *data): # Compute Jacobian jacobian = torch.func.vmap(torch.func.jacrev(self.repr_forward), randomness="different")(flat_data) - jacobian = jacobian.squeeze() + # Squeeze singleton axes but never the batch axis: with batch_size == 1 a + # bare squeeze() drops it and the downstream indexing raises IndexError. + for axis in reversed([i for i, size in enumerate(jacobian.shape) if size == 1 and i != 0]): + jacobian = jacobian.squeeze(axis) neuron_attr = jacobian.swapaxes(0, 1) h_dim = neuron_attr.shape[0] diff --git a/deeptab/distributions/categorical.py b/deeptab/distributions/categorical.py index 618da577..5cb8d024 100644 --- a/deeptab/distributions/categorical.py +++ b/deeptab/distributions/categorical.py @@ -114,6 +114,16 @@ def __init__( self.total_count = total_count self.probs_transform = self.get_transform(prob_transform) + def forward(self, predictions): + """Apply the softmax across the class axis. + + The base implementation looks for a ``_transform`` attribute + per parameter, but this family's parameters are per-class (``p_0``, + ``p_1``, ...) and share a single transform over the whole row, so + without this override ``predict(raw=False)`` returned raw logits. + """ + return self.probs_transform(predictions) + def compute_loss(self, predictions, y_true): probs = self.probs_transform(predictions) diff --git a/deeptab/nn/blocks/common.py b/deeptab/nn/blocks/common.py index 0ada7006..40fb5a3c 100644 --- a/deeptab/nn/blocks/common.py +++ b/deeptab/nn/blocks/common.py @@ -109,7 +109,9 @@ def backward(ctx, grad_output): # type: ignore grad_input = grad_output.clone() grad_input[output == 0] = 0 - v_hat = grad_input.sum(dim=dim) / supp_size.to(output.dtype).squeeze() + # squeeze(dim), not squeeze(): a bare squeeze also drops unrelated + # size-1 axes, which then broadcast into silently wrong gradients. + v_hat = grad_input.sum(dim=dim) / supp_size.to(output.dtype).squeeze(dim) v_hat = v_hat.unsqueeze(dim) grad_input = torch.where(output != 0, grad_input - v_hat, grad_input) return grad_input, None diff --git a/tests/test_distribution_misc_fixes.py b/tests/test_distribution_misc_fixes.py new file mode 100644 index 00000000..00315522 --- /dev/null +++ b/tests/test_distribution_misc_fixes.py @@ -0,0 +1,76 @@ +"""Regression tests for three independent numerical/robustness defects.""" + +from typing import cast + +import numpy as np +import pandas as pd +import pytest +import torch + +from deeptab.distributions import MultinomialDistribution +from deeptab.nn.blocks.common import sparsemax + +QUIET = { + "accelerator": "cpu", + "devices": 1, + "enable_progress_bar": False, + "enable_model_summary": False, + "logger": False, +} + + +class TestMultinomialProbabilities: + def test_forward_returns_probabilities_not_logits(self): + """predict(raw=False) must apply the softmax the docstring promises.""" + family = MultinomialDistribution(num_classes=3) + logits = torch.tensor([[2.0, 0.0, -1.0], [0.0, 0.0, 0.0]]) + + probs = family(logits) + assert torch.allclose(probs, torch.softmax(logits, dim=-1), atol=1e-6) + assert torch.allclose(probs.sum(dim=-1), torch.ones(2), atol=1e-6) + assert (probs >= 0).all() + + +class TestSparsemaxGradients: + def test_gradient_unaffected_by_a_leading_size_one_axis(self): + """A bare squeeze() in backward dropped unrelated size-1 axes.""" + data = torch.tensor([[1.0, 0.5, -0.2, 2.0, 0.1]]) + + flat = data.clone().requires_grad_(True) + cast(torch.Tensor, sparsemax(flat, dim=-1)).sum().backward() + + nested = data.clone().unsqueeze(1).requires_grad_(True) # shape (1, 1, 5) + cast(torch.Tensor, sparsemax(nested, dim=-1)).sum().backward() + + assert flat.grad is not None and nested.grad is not None + torch.testing.assert_close(nested.grad.squeeze(1), flat.grad) + + def test_gradient_matches_batched_reference(self): + rows = torch.tensor([[1.0, 0.5, -0.2, 2.0, 0.1], [0.3, 0.3, 0.9, -1.0, 0.2]]) + + batched = rows.clone().requires_grad_(True) + cast(torch.Tensor, sparsemax(batched, dim=-1)).sum().backward() + + single = rows[:1].clone().requires_grad_(True) + cast(torch.Tensor, sparsemax(single, dim=-1)).sum().backward() + + assert batched.grad is not None and single.grad is not None + torch.testing.assert_close(single.grad[0], batched.grad[0]) + + +class TestTangosSingleRowBatch: + def test_fit_with_a_trailing_single_row_batch(self): + """An odd-sized dataset leaves a batch of one; that must not crash.""" + pytest.importorskip("torch.func") + from deeptab.models.experimental import TangosRegressor + + rng = np.random.default_rng(0) + # 42 rows with the default 0.2 split leaves 33 training rows, so the + # final batch of 16 holds exactly one row. + n = 42 + X = pd.DataFrame({"a": rng.normal(size=n), "b": rng.normal(size=n)}) + y = rng.normal(size=n) + + model = TangosRegressor() + model.fit(X, y, max_epochs=1, batch_size=16, random_state=101, **QUIET) + assert np.isfinite(model.predict(X)).all()