flax

0734d006 - [linen] Only `pmean` the mean if it is non-zero, in normalization modules.