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: 633689990

Author: Jake VanderPlas

language_table/train/policy.py

@@ -67,7 +67,7 @@ def __init__(self, time_step_spec, action_spec, model, checkpoint_path,
 
   def _run_action_inference(self, observation):
     # Add a batch dim.
-    observation = jax.tree_map(lambda x: jnp.expand_dims(x, 0), observation)
+    observation = jax.tree.map(lambda x: jnp.expand_dims(x, 0), observation)
 
     normalized_action = self.model.apply(
         self.variables, observation, train=False)

language_table/train/train.py

@@ -44,7 +44,7 @@ def _body_fun(step, state_and_metrics):
     train_rng = jax.random.fold_in(train_rng, jax.lax.axis_index("batch"))
     new_state, metrics_update = agent.train(
         state=state,
-        batch=jax.tree_map(lambda x: x[step], batches),
+        batch=jax.tree.map(lambda x: x[step], batches),
         rng=train_rng,
     )
     return new_state, metrics_update
@@ -101,7 +101,7 @@ def train(
   )
   train_iter = train_ds.as_numpy_iterator()
   rng, agent_rng = jax.random.split(rng)
-  sample_batch = jax.tree_map(lambda x: x[0][0], next(train_iter))
+  sample_batch = jax.tree.map(lambda x: x[0][0], next(train_iter))
   agent = create_agent(
       config.agent_name,
       config.model_name,
@@ -257,7 +257,7 @@ def create_agent(
 
 def merge_batch_stats(replicated_state):
   """Merge model batch stats."""
-  if jax.tree_leaves(replicated_state.batch_stats):
+  if jax.tree.leaves(replicated_state.batch_stats):
     cross_replica_mean = jax.pmap(lambda x: jax.lax.pmean(x, "x"), "x")
     return replicated_state.replace(
         batch_stats=cross_replica_mean(replicated_state.batch_stats)