fix omnistaging issues

[jheek]

These changes are required to be forward compatible with the upcoming omnistaging feature in Jax (jax-ml/jax#3370).

[codecov-commenter]

Codecov Report

Merging #335 into master will decrease coverage by 0.28%.
The diff coverage is 92.30%.

@@            Coverage Diff             @@
##           master     #335      +/-   ##
==========================================
- Coverage   79.32%   79.03%   -0.29%     
==========================================
  Files          34       34              
  Lines        2254     2247       -7     
==========================================
- Hits         1788     1776      -12     
- Misses        466      471       +5     

Impacted Files	Coverage Δ
flax/jax_utils.py	62.37% <91.66%> (-2.05%)	⬇️
flax/nn/utils.py	93.33% <100.00%> (-4.67%)	⬇️
flax/nn/stochastic.py	71.05% <0.00%> (-2.64%)	⬇️
flax/nn/base.py	91.70% <0.00%> (-0.23%)	⬇️
flax/nn/normalization.py	82.71% <0.00%> (+0.21%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 484cc11...ee20ca2. Read the comment docs.

[levskaya]

LGTM - thanks for fixing these for the upcoming changes!