Fix RMSNorm scale zero-initialization

[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.