Fix jax.tree_map deprecation warnings.

PiperOrigin-RevId: 622995953
google-deepmind · Apr 9, 2024 · eff8b6d · eff8b6d
1 parent aea1d5f
commit eff8b6d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optax/tree_utils/_state_utils.py b/optax/tree_utils/_state_utils.py
@@ -149,13 +149,13 @@ def tree_map_params(
 
   def map_params(maybe_placeholder_value, value):
     if isinstance(maybe_placeholder_value, _ParamsPlaceholder):
-      return jax.tree_map(f, value, *rest, is_leaf=is_leaf)
+      return jax.tree_util.tree_map(f, value, *rest, is_leaf=is_leaf)
     elif transform_non_params is not None:
       return transform_non_params(value)
     else:
       return value
 
-  return jax.tree_map(
+  return jax.tree_util.tree_map(
       map_params,
       state_with_placeholders,
       state,