fix: Multinomial probabilities, sparsemax gradients, Tangos single-row batches

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]

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 ``<param_name>_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)
 

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

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()