cleanup: replace deprecated jax.tree_map with jax.tree_util.tree_map

[copybara-service]

cleanup: replace deprecated jax.tree_map with jax.tree_util.tree_map

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov-commenter]

Codecov Report

Merging #2311 (8d27c74) into main (80cb1b2) will not change coverage.
The diff coverage is 72.09%.

@@           Coverage Diff           @@
##             main    #2311   +/-   ##
=======================================
  Coverage   75.83%   75.83%           
=======================================
  Files          61       61           
  Lines        5383     5383           
=======================================
  Hits         4082     4082           
  Misses       1301     1301           

Impacted Files	Coverage Δ
flax/linen/summary.py	93.75% <ø> (ø)
flax/optim/base.py	34.10% <0.00%> (ø)
flax/optim/dynamic_scale.py	0.00% <0.00%> (ø)
flax/optim/weight_norm.py	27.84% <0.00%> (ø)
flax/core/axes_scan.py	95.45% <100.00%> (ø)
flax/core/lift.py	95.91% <100.00%> (ø)
flax/core/scope.py	89.88% <100.00%> (ø)
flax/linen/partitioning.py	82.32% <100.00%> (ø)
flax/linen/transforms.py	89.68% <100.00%> (ø)