flax

[linen] Add alternative, more numerically stable, variance calculation to `LayerNorm`.

#3194

Merged