[meemoprasad]

RMSNorm.__init__ in src/timesfm/flax/normalization.py initializes
self.scale with jnp.zeros instead of jnp.ones:

self.scale = nnx.Param(jnp.zeros(shape=(num_features,)))

Since __call__ does normed_inputs *= self.scale, this makes
RMSNorm output all zeros on a fresh forward pass, before any
checkpoint load overwrites the parameter.

LayerNorm in the same file correctly initializes its scale with
jnp.ones — this looks like a copy-paste bug between the two classes.

Currently masked in inference when checkpoints are always loaded
before a forward pass, but affects:

• standalone instantiation of RMSNorm (e.g. in tests)
• any future from-scratch training / fine-tuning of the flax backend

Fix: jnp.zeros → jnp.ones, matching LayerNorm's pattern.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[meemoprasad]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

CLA is now signed and verified — ready for review whenever a maintainer gets a chance.

[sylvesterkaczmarek]

This fixes the zero-initialized scale, but there is no regression covering the failure mode. With the old initializer, a fresh RMSNorm maps every input to zero because the normalized activations are multiplied by an all-zero scale, and the current suite would not catch that coming back. Could we add a small test that a newly constructed layer has unit scale and preserves nonzero normalized output?

[meemoprasad]

Thanks for the review. Before adding a test, I want to flag something I missed: tests/test_torch_layers.py:160-171 (test_zero_scale_produces_zeros) documents zero-initialised scale as an intentional property (“relying on the residual connection to carry signal”), and the torch RMSNorm (src/timesfm/torch/normalization.py:31) also initialises with zeros. So changing only the flax backend to ones would make the two backends diverge.

Maintainers: is zero-init intended for both backends? If yes, I’ll close this PR or turn it into a docstring note on RMSNorm explaining the zero-init. If you’d prefer unit init, I can update both backends and the existing test in this PR.

[sylvesterkaczmarek]

Good catch. test_zero_scale_produces_zeros explicitly documents zero-initialized RMSNorm as intentional, and the Torch backend uses the same default. My earlier request for a unit-scale regression assumed the initialization was a bug. Please disregard that request; the backends should stay consistent unless maintainers decide to change the design.

[meemoprasad]

Thanks @sylvesterkaczmarek, agreed. I’ll repurpose this PR: revert the init change and instead add a short docstring to RMSNorm in both the flax and torch backends explaining that zero init is intentional, so the next reader doesn’t hit the same question. Happy to close it instead if maintainers prefer.

[meemoprasad]

Reverted the init change, kept it as zeros. Added a docstring on RMSNorm in both backends explaining the intentional zero-init, pointing to test_zero_scale_produces_zeros. Ready for review as a docs-only PR.

[sylvesterkaczmarek]

Rechecked head 81a1b77. The initialization is unchanged in both backends, and the new RMSNorm docstrings accurately document the intended zero scale and point to the existing Torch regression. This resolves my earlier concern. Looks good as a documentation-only change.