Add mean_init and var_init parameters to the BatchNorm module

[matthias-wright]

Hi all,

Adding mean_init and var_init parameters to the BatchNorm module would make it easier to load checkpoints from other frameworks. I am currently doing this but the mean_init and var_init parameters would make it more convenient, especially for training / fine-tuning.

These parameters could work exactly like bias_init and scale_init.

Inserting these two lines here:

mean_init: Callable[[Shape], Array] = lambda s: jnp.zeros(s, jnp.float32)
var_init: Callable[[Shape], Array] = lambda s: jnp.ones(s, jnp.float32)

and then replacing the lambdas here with the self.mean_init and self.var_init would make this a non-breaking change.

I am happy to submit a PR if there is interest. What do you think?

Cheers!