Add use_fast_variance option to GroupNorm and BatchNorm to allow disabling it.

[copybara-service]

Add use_fast_variance option to GroupNorm and BatchNorm to allow disabling it.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[codecov-commenter]

Codecov Report

Merging #3256 (bc6a6c1) into main (8da8c46) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #3256   +/-   ##
=======================================
  Coverage   82.34%   82.34%           
=======================================
  Files          54       54           
  Lines        6088     6090    +2     
=======================================
+ Hits         5013     5015    +2     
  Misses       1075     1075           

Files Changed	Coverage Δ
flax/linen/normalization.py	97.46% <100.00%> (+0.03%)	⬆️