Why use standard LayerNorm over RMSNorm in JAX/Flax? Does XLA fusion negate the algorithmic disadvantage?

[AnirudhSKrishnan]

I'm implementing normalization in Flax and trying to decide between standard LayerNorm and RMSNorm.

Usually RMSNorm is preferred because it's faster (skips the mean calculation/centering). But with jax.jit, doesn't XLA fuse the standard LayerNorm math into a single kernel anyway?

Does the "two-pass" nature of standard LayerNorm actually matter once it's compiled, or is the performance difference negligible compared to the memory bandwidth savings from fusion? Just trying to figure out if it's worth writing a custom RMSNorm or if I should just stick to the standard nn.LayerNorm.

Thanks.

[vfdev-5]

Flax NNX also provides nnx.RMSNorm