Hi, @ptrblck
A similar question about Jax and Pytorch.
For BatchNorm, I noticed that there is a difference between Jax 1 and Pytorch.
I.e.
Jax computation order:  (inputs - mean) *  (scale * rsqrt(var + eps))
Pytorch:  (inputs - mean) * scale * rsqrt(var + eps).
I know they are theoretically equivalent.
Do you think this computation order would make difference?
Thanks in advance.