Skip to content

fix(core): second-pass review findings - #439

Open
ChrisW09 wants to merge 2 commits into
mainfrom
fix/second-pass-findings
Open

ChrisW09 wants to merge 2 commits into
mainfrom
fix/second-pass-findings

Conversation

@ChrisW09

Copy link
Copy Markdown
Collaborator

Fixes #426 (items 1, 2, 4 and 5; the CRPS-fallback sigma, the MambaOriginal cache method and the TabTransformer zero-numeric-feature guard are left for separate PRs since each needs a design call).

Problems fixed

  1. BaseModel.encode() stray call. The grad=False branch (the default, used by the public clf.encode(X) mixin) hardcoded embeddings = self.encoder(x) before overwriting it with layer(x). Mambular / MambAttention / MambaTab expose self.mamba and no encoder, so encode() raised AttributeError; every other architecture ran its encoder twice and threw the first result away. The symmetric grad=True branch never had the extra line.
  2. Phantom family keys in _validate_fit_inputs. The strictly-positive check gated on "inversegaussian", which does not exist in DISTRIBUTION_REGISTRY (the real key is "inversegamma"), so it never fired — and lognormal needs the same guard.
  3. sparsemax mutated its input in place. input_ -= max_val inside Function.forward; ODST (nn/blocks/node.py) passes its feature_selection_logits Parameter directly, so the stored values were shifted by -max on first use. The output is shift-invariant so predictions were unaffected, but weight decay and saved checkpoints operated on corrupted values.
  4. TaskModel.train_features was never initialised. The step methods guard with ... and self.train_features is not None, which presumes a None default, but the attribute is only created in setup(stage="fit") — so validate/predict on a candidate model restored from a checkpoint raised AttributeError from the guard itself.

Tests

New tests/test_second_pass_fixes.py: sparsemax leaves its input untouched (and still returns a valid simplex projection), the positivity check fires for all three strictly-positive families and the checked keys actually exist in the registry, and TaskModel starts with train_features is None. tests/test_exceptions.py, tests/test_nn_blocks.py and tests/test_models.py pass unchanged (424 passed).

🤖 Generated with Claude Code

ChrisW09 and others added 2 commits July 27, 2026 21:34
- BaseModel.encode(): drop the stray self.encoder(x) call in the no-grad
  branch. Mamba architectures have no 'encoder' attribute, so the public
  clf.encode(X) path raised AttributeError for them; for every other
  model it ran the encoder twice and discarded the first result. The
  grad=True branch never had the extra call.
- _validate_fit_inputs gated the strictly-positive check on
  'inversegaussian', which is not in DISTRIBUTION_REGISTRY (the real key
  is 'inversegamma'), so the check never fired; lognormal needs it too.
- sparsemax shifted its input in place. ODST passes the
  feature_selection_logits Parameter directly, so the stored parameter
  values were corrupted by -max on first use. Output is shift-invariant,
  so predictions were unaffected, but weight decay and checkpoints saw
  the shifted values.
- TaskModel now initialises train_features/train_targets to None, so the
  'is not None' guards in the step methods work for candidate models
  restored from a checkpoint without an in-process fit.

Fixes #426

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sparsemax is typed as returning Tensor | None via the autograd Function.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant