[BUG] Ensemble LSS models (TabM, Trompt) crash for every single-parameter distribution family

[ChrisW09]

Describe the bug

Ensemble architectures unconditionally squeeze(-1) their output. When the distribution family has
exactly one parameter that squeeze deletes the parameter axis, so TabMLSS / TromptLSS raise for
poisson, tweedie, dirichlet and categorical.

Reproduced: TabMLSS(...).fit(X, y, family="poisson") -> IndexError: too many indices for tensor of dimension 1, raised from deeptab/distributions/poisson.py:29
(predictions[:, self.param_names.index("rate")]).

Shape proof: Trompt(..., num_classes=1) returns (5, 6) where num_classes=2 returns (5, 6, 2) —
the trailing parameter axis is dropped only in the one-parameter case.

---

Ensemble LSS models crash for every single-parameter distribution family

Where: deeptab/architectures/tabm.py (206-207 (also deeptab/architectures/experimental/trompt.py:88, consumed at deeptab/training/lightning_module.py:329-335 and deeptab/models/lss_base.py:374-378))

Ensemble architectures unconditionally squeeze(-1) their output, which deletes the distribution-parameter axis when the family has exactly one parameter, so TabMLSS/TromptLSS raise IndexError for poisson, tweedie, dirichlet and categorical.

Observed: IndexError: too many indices for tensor of dimension 1 raised from deeptab/distributions/poisson.py:29 (predictions[:, self.param_names.index('rate')]) via lightning_module.py:332. Trompt forward with num_classes=1 returns shape (5, 6) instead of (5, 6, 1); TabM does the same at tabm.py:207 (average_ensembles defaults to False, so returns_ensemble=True and the squeeze is the default path). Same crash for family='tweedie'. If training somehow succeeded, SklearnBaseLSS.predict would fail identically: predictions.mean(dim=1) yields a 1-D tensor that family(...) then indexes with [:, idx].

Expected: Fit and predict should succeed for one-parameter families exactly as they do for 'normal' (2 params) — the ensemble output should keep its trailing parameter axis of size param_count, i.e. squeeze(-1) should be applied only when the model is not in LSS mode (or only when num_classes == 1 AND lss is False).

Repro

import warnings, os; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models import TabMLSS
from deeptab.models.experimental import TromptLSS
from deeptab.configs import TabMConfig, TromptConfig, TrainerConfig

rng = np.random.default_rng(0); N = 48
X = pd.DataFrame({'a': rng.normal(size=N), 'b': rng.normal(size=N)})
y = np.abs(rng.normal(size=N)).round()          # non-negative counts
kw = dict(accelerator='cpu', devices=1, enable_progress_bar=False,
          enable_model_summary=False, logger=False)

for cls, cfg in [(TabMLSS, TabMConfig()), (TromptLSS, TromptConfig())]:
    m = cls(model_config=cfg, trainer_config=TrainerConfig(max_epochs=1, batch_size=16))
    m.fit(X, y, family='poisson', **kw)          # -> IndexError

# direct shape proof:
import torch
from deeptab.architectures.experimental.trompt import Trompt
ni = {'a': {'dimension': 4, 'preprocessing': None, 'categories': None}}
print(Trompt((ni, {}, {}), num_classes=1, config=TromptConfig())([torch.randn(5,4)], [], []).shape)  # (5, 6)
print(Trompt((ni, {}, {}), num_classes=2, config=TromptConfig())([torch.randn(5,4)], [], []).shape)  # (5, 6, 2)

---

Expected behavior
One-parameter families should fit and predict exactly as normal (2 params) does; the ensemble output
must keep its trailing parameter axis of size param_count. The squeeze(-1) should apply only when
the model is not in LSS mode.

Screenshots
n/a

Desktop (please complete the following information):

• OS: macOS (Darwin 25.5.0, arm64)
• Python version: 3.11.15
• deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.