Restore previous _compute_stats behavior in local case.

[copybara-service]

Restore previous _compute_stats behavior in local case.

Some earlier code refactors introduced a concat+split "no-op" to the non-distributed codepath
that turns out not to be such a no-op and can cause issues in some compiled code. We polish
the refactor a bit to avoid this introduced change.

[codecov-commenter]

Codecov Report

Merging #3266 (e6d0e97) into main (fe54d39) will decrease coverage by 0.01%.
The diff coverage is 91.66%.

@@            Coverage Diff             @@
##             main    #3266      +/-   ##
==========================================
- Coverage   82.35%   82.34%   -0.01%     
==========================================
  Files          54       54              
  Lines        6092     6095       +3     
==========================================
+ Hits         5017     5019       +2     
- Misses       1075     1076       +1     

Files Changed	Coverage Δ
flax/linen/normalization.py	96.89% <91.66%> (-0.58%)	⬇️