[linen] Why not handle batchnorm immutably?

[MRiabov]

Hello Flax community,
Why can't batchnorm be handled immutably? it occurs to me it can easily be done via...

1. Implicitly initializing a batch_stats array with leading size set by the user (e.g. 200 states)
2. During majority of training, take an average from that array in a jit manner,
3. In the beginning, when some values are zeroes...

last_item = jax.nonzero(batch_stats)[-1]
norm = mean(batch_stats[:last_item])
#and etc logic implementing getting values before nonzero.

Am I missing something? Definitely beats having the user write all the extra code for mutable state.

#note: I'm somewhat new to flax and ML, so I may be wrong.