Replace deprecated jax.tree_* functions with jax.tree.*

The top-level `jax.tree_*` aliases have long been deprecated, and will soon be removed. Alternate APIs are in `jax.tree_util`, with shorter aliases in the `jax.tree` submodule, added in JAX version 0.4.25. PiperOrigin-RevId: 633664602
google · May 14, 2024 · f0f8df1 · f0f8df1
1 parent b172f5a
commit f0f8df1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -117,6 +117,6 @@ Last release before dropping support for Python 3.8 and 3.9
 
 ## v0.0.12
 
--   Switch from `jax.tree_map` (deprecated since JAX 0.4.26) to
+-   Switch from `jax.tree.map` (deprecated since JAX 0.4.26) to
     `jax.tree_util.tree_map`.
 -   Improvements to parameter overview.