fix: replace jax.tree_* with jax.tree

Author: epignatelli

examples/purejaxrl/ppo_minigrid.py

@@ -225,14 +225,14 @@ def _loss_fn(params, traj_batch, gae, targets):
                 ), "batch size must be equal to number of steps * number of envs"
                 permutation = jax.random.permutation(_rng, batch_size)
                 batch = (traj_batch, advantages, targets)
-                batch = jax.tree_util.tree_map(
+                batch = jax.tree.map(
                     lambda x: x.reshape((batch_size,) + x.shape[2:]), batch
                 )
-                shuffled_batch = jax.tree_util.tree_map(
+                shuffled_batch = jax.tree.map(
                     lambda x: jnp.take(x, permutation, axis=0), batch
                 )
                 # Mini-batch Updates
-                minibatches = jax.tree_util.tree_map(
+                minibatches = jax.tree.map(
                     lambda x: jnp.reshape(
                         x, [config["NUM_MINIBATCHES"], -1] + list(x.shape[1:])
                     ),

navix/agents/ppo.py

@@ -252,15 +252,15 @@ def update(self, train_state: TrainingState, _) -> Tuple[TrainingState, Dict]:
             ), "batch size must be equal to number of steps * number of envs"
             permutation = jax.random.permutation(rng_1, n_samples)
             samples = (experience, advantages, targets, values)  # (T, N, ...)
-            samples = jax.tree_util.tree_map(
+            samples = jax.tree.map(
                 lambda x: x.reshape((n_samples,) + x.shape[2:]), samples
             )  # (T * N, ...)
-            shuffled_batch = jax.tree_util.tree_map(
+            shuffled_batch = jax.tree.map(
                 lambda x: jnp.take(x, permutation, axis=0), samples
             )  # (T * N, ...)
 
             # One epoch update over all mini-batches
-            minibatches = jax.tree_util.tree_map(
+            minibatches = jax.tree.map(
                 lambda x: jnp.reshape(
                     x, (self.hparams.num_minibatches, -1) + tuple(x.shape[1:])
                 ),

navix/entities.py

@@ -66,7 +66,7 @@ class Entity(Positionable, HasTag, HasSprite):
     To create an entity, use the `create` method."""
 
     def __getitem__(self: T, idx) -> T:
-        return jax.tree_util.tree_map(lambda x: x[idx], self)
+        return jax.tree.map(lambda x: x[idx], self)
 
     @property
     def name(self) -> str:

navix/environments/key_corridor.py

@@ -24,7 +24,6 @@
 import jax
 import jax.numpy as jnp
 from jax import Array
-import jax.tree_util as jtu
 
 from navix import observations, rewards, terminations
 
@@ -116,7 +115,7 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
                     open=jnp.asarray(0),
                 )
             )
-        doors = jtu.tree_map(lambda *x: jnp.stack(x), *doors)
+        doors = jax.tree.map(lambda *x: jnp.stack(x), *doors)
 
         entities = {
             "player": player[None],

navix/experiment.py

@@ -171,10 +171,10 @@ def search(hparam_set_sample):
         # average over seeds
         for i in range(len_search_set):
             print("Logging results for hparam set:", search_set)
-            hparams = jax.tree_map(lambda x: x[i], search_set)
+            hparams = jax.tree.map(lambda x: x[i], search_set)
             config = {**vars(self), **asdict(hparams)}
             wandb.init(project=self.name, config=config, group=self.group)
-            log = jax.tree_map(lambda x: jnp.mean(x[i], axis=0), logs)
+            log = jax.tree.map(lambda x: jnp.mean(x[i], axis=0), logs)
             self.agent.log_on_train_end(log)
             wandb.finish()
         logging_time = time.time() - start_time

navix/grid.py

@@ -26,7 +26,6 @@
 import jax
 import jax.numpy as jnp
 from jax import Array
-import jax.tree_util as jtu
 from flax import struct