Migrate from jax.tree_util legacy APIs to new jax.tree API.#1066

Merged

copybara-service[bot] merged 3 commits intogoogle-deepmind:mainfrom carlosgmartin:jax_tree_util_legacy

Sep 23, 2024

+331-342

Commits on Sep 20, 2024