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

[ChrisW09]

Fixes #453 — three independent defects, each with its own regression test.

1. MultinomialDistribution.forward returned raw logits

BaseDistribution.forward looks up a <param_name>_transform attribute for each parameter. This family names its parameters per class (p_0, p_1, …) and keeps a single shared probs_transform, so every lookup missed and silently fell back to identity:

MultinomialDistribution(num_classes=3)(torch.tensor([[2.0, 0.0, -1.0]]))
# -> [[2.0, 0.0, -1.0]]   (logits, not probabilities)

compute_loss applies the softmax correctly, so only predict(raw=False) was affected — it returned logits labelled as probabilities. Fixed with a small forward override.

2. sparsemax backward produced silently wrong gradients

grad_input.sum(dim=dim) / supp_size.to(output.dtype).squeeze() — a bare squeeze() also drops unrelated size-1 axes, and the resulting shape mismatch broadcasts into wrong gradient values whenever a non-reduced dimension has size 1 (e.g. NODE with depth 1). Now squeeze(dim), which removes only the reduced axis.

3. Tangos crashed on single-row batches

penalty_forward calls jacobian.squeeze(), which removes the batch axis when the batch holds one row — routine for the final batch of an odd-sized dataset. The following neuron_attr.shape[2] then raised IndexError. Singleton axes other than the batch axis are now squeezed individually.

Tests

New tests/test_distribution_misc_fixes.py: Multinomial forward matches torch.softmax and sums to 1; sparsemax gradients are unchanged by adding a leading size-1 axis and match a batched reference; and TangosRegressor fits a 42-row dataset whose default split leaves a final batch of exactly one row. Verified 2 of the 4 fail on main (the sparsemax pair needs the size-1-axis case to differ, which the first test pins).

🤖 Generated with Claude Code