flax
0797d053 - Use `jax.tree_util.tree_map` instead of deprecated `jax.tree_map`.

Commit

1 year ago

Use `jax.tree_util.tree_map` instead of deprecated `jax.tree_map`. PiperOrigin-RevId: 610521805

References

Author

IvyZX

Committer

a-googler

Parents