fix(architectures): keep the distribution-parameter axis in ensemble LSS models - #457
Open
ChrisW09 wants to merge 1 commit into
Open
fix(architectures): keep the distribution-parameter axis in ensemble LSS models#457ChrisW09 wants to merge 1 commit into
ChrisW09 wants to merge 1 commit into
Conversation
…LSS models
TabM and Trompt unconditionally squeeze the trailing axis of their
ensemble output. That is correct for single-output regression and binary
classification, where the axis has size 1 and the loss expects (B, E) --
but in LSS mode the trailing axis is the distribution-parameter axis. For
a family with exactly one parameter it also has size 1, so it was
deleted, and compute_loss's predictions[:, member] then produced a 1-D
tensor that the distribution indexed with [:, idx]:
IndexError: too many indices for tensor of dimension 1
raised from deeptab/distributions/poisson.py:29 for TabMLSS/TromptLSS
with family in {poisson, tweedie, dirichlet, categorical}. Multi-parameter
families were unaffected because the squeeze was a no-op for them.
The squeeze is now skipped in LSS mode, which is a no-op for the families
that already worked.
Fixes #447
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #447
Problem
TabMandTromptunconditionallysqueeze(-1)their ensemble output. That is right forsingle-output regression and binary classification — the trailing axis has size 1 and the ensemble loss
expects
(B, E). But in LSS mode the trailing axis is the distribution-parameter axis, and for afamily with exactly one parameter it also has size 1, so it was deleted.
compute_lossthen doespredictions[:, ensemble_member], yielding a 1-D tensor that the distributionindexes with
[:, idx]:Reproduced for
TabMLSSandTromptLSSwithfamilyin{poisson, tweedie, dirichlet, categorical}.Multi-parameter families were unaffected because the squeeze was a no-op for them.
Shape proof:
Trompt(..., num_classes=1)returned(5, 6)wherenum_classes=2returned(5, 6, 2).Fix
Skip the squeeze in LSS mode (
self.hparams.lss, already threaded into every architecture byTaskModel). For the multi-parameter families that already worked this changes nothing, since thesqueeze was a no-op there.
Tests
New
tests/test_ensemble_lss.py:TabMLSSfits and predicts finitely for bothpoisson(1 param) andnormal(2 params), plus a direct shape assertion that the LSS ensemble output keeps its parameteraxis while the non-LSS one still collapses. Verified 2 of the 3 fail on
main.tests/test_models.pyand
tests/test_lss_base.pypass unchanged.🤖 Generated with Claude Code