Use jax.tree_util to avoid deprecation warnings.

PiperOrigin-RevId: 463276750

Author: fidlej

mctx/_src/search.py

@@ -217,7 +217,7 @@ def expand(
   chex.assert_shape([parent_index, action, next_node_index], (batch_size,))
 
   # Retrieve states for nodes to be evaluated.
-  embedding = jax.tree_map(
+  embedding = jax.tree_util.tree_map(
       lambda x: x[batch_range, parent_index], tree.embeddings)
 
   # Evaluate and create a new node.
@@ -333,7 +333,7 @@ def update_tree_node(
           tree.node_values, value, node_index),
       node_visits=batch_update(
           tree.node_visits, new_visit, node_index),
-      embeddings=jax.tree_map(
+      embeddings=jax.tree_util.tree_map(
           lambda t, s: batch_update(t, s, node_index),
           tree.embeddings, embedding))
 
@@ -373,7 +373,7 @@ def _zeros(x):
       children_visits=jnp.zeros(batch_node_action, dtype=jnp.int32),
       children_rewards=jnp.zeros(batch_node_action, dtype=data_dtype),
       children_discounts=jnp.zeros(batch_node_action, dtype=data_dtype),
-      embeddings=jax.tree_map(_zeros, root.embedding),
+      embeddings=jax.tree_util.tree_map(_zeros, root.embedding),
       root_invalid_actions=root_invalid_actions,
       extra_data=extra_data)
 

mctx/_src/tree.py

@@ -117,7 +117,7 @@ def infer_batch_size(tree: Tree) -> int:
   """Recovers batch size from `Tree` data structure."""
   if tree.node_values.ndim != 2:
     raise ValueError("Input tree is not batched.")
-  chex.assert_equal_shape_prefix(jax.tree_leaves(tree), 1)
+  chex.assert_equal_shape_prefix(jax.tree_util.tree_leaves(tree), 1)
   return tree.node_values.shape[0]