Bug Report

GPJax version: 0.18.0

Current behaviour

The Zero mean function is trainable, and drifts away from zero during optimisation. Its own
docstring says otherwise:

Unlike the Constant mean function, the constant scalar zero is fixed, and cannot be treated
as a model hyperparameter and learned during training.

before fit, Zero constant = 0.0
after  fit, Zero constant = 5.088802112651860

This is a regression of #330, which was fixed in #500.

Expected behaviour

Zero() returns zero for all inputs, before and after fitting. This matches the docstring and
matches GPflow.

Steps to reproduce

import jax; jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp, optax as ox, gpjax as gpx

x = jnp.linspace(0, 1, 50, dtype=jnp.float64).reshape(-1, 1)
y = 5.0 + jnp.sin(4 * x)                      # data with a large non-zero mean
D = gpx.Dataset(X=x, y=y)

prior = gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF())
posterior = prior * gpx.likelihoods.Gaussian(num_datapoints=D.n)

opt, _ = gpx.fit(
    model=posterior,
    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
    train_data=D, optim=ox.adam(0.05), num_iters=300,
    key=jax.random.PRNGKey(0), verbose=False,
)
print(opt.prior.mean_function.constant)   # 5.0888..., expected 0.0

Also reproduces on the state-space path: d(state_space_mll)/d(constant) = 0.0366.

Root cause

Not the missing wrapper on its own — the trainability contract changed underneath it.

Under nnx, fit optimised only Parameter instances (trainable: nnx.filterlib.Filter = Parameter), so Zero's raw array was inert by construction — which is why #530 could drop
Static and remain correct. Under Equinox, fit optimises every array leaf
(fit.py:222, eqx.partition(model, eqx.is_array)), so the same raw array became trainable.

Zero was relying on "raw array = frozen", and that stopped being true.

The test regression

A guard existed and was weakened rather than fixed. #614 deleted test_zero_mean_remains_zero
(which fitted a dataset with mean 50 and asserted the mean stayed at zero) and replaced it with
test_zero_mean_initialises_at_zero, whose docstring reads:

Note: with equinox, the constant is a plain array leaf and is trainable. The invariant we
test is that the initial value is zero, not that it stays zero after optimisation.

Separately, test_zero_mean_function_uses_raw_value
(tests/test_mean_functions.py:290) asserts not isinstance(meanf.constant, AbstractUnwrappable) — it now pins the defect in place.

So the behaviour change was observed and rationalised in the test suite, not missed. Worth
noting for its own sake: a test that is relaxed to accommodate a migration is exactly the
signal we most need to not lose.

Fix

One line, verified to hold at 0.0 across the same 300-step fit:

class Zero(Constant):
    def __init__(self):
        super().__init__(constant=paramax.non_trainable(jnp.array(0.0)))

Alongside it:

• restore test_zero_mean_remains_zero as an end-to-end fit assertion;
• invert test_zero_mean_function_uses_raw_value, which currently asserts the bug.

Scope

I swept for other things relying on "raw array = frozen" across the nnx→Equinox boundary.
Zero is the only one: GraphKernel already uses paramax.non_trainable for its Laplacian
and eigendecomposition, and no other docstring claims a fixed parameter.

White.lengthscale is a related but distinct problem — a phantom parameter with identically
zero gradient, so it pollutes summary() and MCMC dimensionality without changing results. It
belongs with the structural discussion in #695, not here.