Skip to content

Make zero-mean fn. constant a Static class - #500

Merged
thomaspinder merged 7 commits into
mainfrom
zero-mean-fix
Apr 9, 2025
Merged

thomaspinder merged 7 commits into
mainfrom
zero-mean-fix

Conversation

@thomaspinder

Copy link
Copy Markdown
Collaborator

Checklist

  • I've formatted the new code by running hatch run dev:format before committing.
  • I've added tests for new code.
  • [] I've added docstrings for the new code.

Description

Set the constant field of the zero-mean function to be a Static parameter to prevent its value being changed from zero during training.

Issue Number: #499

@thomaspinder
thomaspinder merged commit 168a106 into main Apr 9, 2025
@thomaspinder
thomaspinder deleted the zero-mean-fix branch April 9, 2025 18:01
thomaspinder added a commit that referenced this pull request Jul 26, 2026
`Zero()` was trainable and drifted towards the data mean during `fit`
(0.0 -> 5.09 on a dataset with mean 5), contradicting its own docstring and
silently changing the posterior mean of every model using the default mean
function. Regression of #330, fixed once in #500.

The cause is a changed trainability contract rather than a lost line. Under
nnx, `fit` optimised only `Parameter` instances, so `Zero`'s bare array was
inert by construction -- which is why #530 could drop the `Static` wrapper
and stay correct. Under Equinox, `fit` partitions on `eqx.is_array`, making
every array leaf trainable, and `Zero` was still relying on the old meaning
of a bare array.

Wrap the constant in `paramax.non_trainable` so the invariant holds by
construction rather than by the ambient filter semantics.

The guard for this was weakened rather than removed: #614 replaced
`test_zero_mean_remains_zero` with a test of the initial value only, and
`test_zero_mean_function_uses_raw_value` asserted the defect outright.
Restore the end-to-end fit assertion and invert the unit test.

Fixes #712


Claude-Session: https://claude.ai/code/session_01Bj9k5fnAZ8JzD4Rg3HMDMj

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant