[ONNX] Reimplement _var_mean to ensure non-negative (#47240)
Summary:
The current `_var_mean` implementation cannot ensure non-negative for variance, because it is actually `E(X^2)-(E(X))^2`: numerically when the dimension number is large and X is close to 0, it can have negative numbers (like our UT shows). The new implementation is `(E(X-E(X))^2)`, it ensures non-negative because the expectation of square is non-negative for sure.
The UT passes for the new implementation (but fails for the existing one). So it is good to go.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/47240
Reviewed By: ejguan
Differential Revision: D24735729
Pulled By: bzinodev
fbshipit-source-id: 136f448dd16622b2b46f40cdf6cb2fccf357c48d