Merge pull request #95 from carlosgmartin:jax_tree_map

PiperOrigin-RevId: 656065747

Author: MctxDev

mctx/_src/policies.py

@@ -385,7 +385,7 @@ def _mask_invalid_actions(logits, invalid_actions):
 
 
 def _get_logits_from_probs(probs):
-  tiny = jnp.finfo(probs).tiny
+  tiny = jnp.finfo(probs.dtype).tiny
   return jnp.log(jnp.maximum(probs, tiny))
 
 

mctx/_src/search.py

@@ -219,7 +219,7 @@ def expand(
   chex.assert_shape([parent_index, action, next_node_index], (batch_size,))
 
   # Retrieve states for nodes to be evaluated.
-  embedding = jax.tree_util.tree_map(
+  embedding = jax.tree.map(
       lambda x: x[batch_range, parent_index], tree.embeddings)
 
   # Evaluate and create a new node.
@@ -335,7 +335,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_util.tree_map(
+      embeddings=jax.tree.map(
           lambda t, s: batch_update(t, s, node_index),
           tree.embeddings, embedding))
 
@@ -375,7 +375,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_util.tree_map(_zeros, root.embedding),
+      embeddings=jax.tree.map(_zeros, root.embedding),
       root_invalid_actions=root_invalid_actions,
       extra_data=extra_data)
 

mctx/_src/tests/policies_test.py

@@ -245,7 +245,7 @@ def test_gumbel_muzero_policy(self):
 
     # Testing max_depth.
     leaf, max_found_depth = _get_deepest_leaf(
-        jax.tree_util.tree_map(lambda x: x[0], policy_output.search_tree),
+        jax.tree.map(lambda x: x[0], policy_output.search_tree),
         policy_output.search_tree.ROOT_INDEX)
     self.assertEqual(max_depth, max_found_depth)
     self.assertEqual(6, policy_output.search_tree.node_visits[0, leaf])