[UPDATE JAX API] Modify all the calls to jax.tree

axlearn/common/causal_lm.py

@@ -199,7 +199,7 @@ def forward(
         live_targets = target_labels >= 0
         num_targets = live_targets.sum()
 
-        logging.info("Module outputs: %s", jax.tree_structure(module_outputs))
+        logging.info("Module outputs: %s", jax.tree_util.tree_structure(module_outputs))
         accumulation = []
         for k, v in flatten_items(module_outputs):
             if re.fullmatch(regex, k):

axlearn/common/gradient_accumulation.py

@@ -33,7 +33,7 @@ def _compute_minibatch_size(input_batch: Nested[Tensor], *, steps: int) -> int:
     if steps <= 0:
         raise ValueError("Accumulation steps need to be a positive integer.")
 
-    input_batch_sizes = jax.tree_leaves(jax.tree.map(lambda x: x.shape[0], input_batch))
+    input_batch_sizes = jax.tree_util.tree_leaves(jax.tree.map(lambda x: x.shape[0], input_batch))
 
     if len(input_batch_sizes) == 0:
         raise ValueError("Input batch is empty.")

axlearn/common/update_transformation.py

@@ -186,7 +186,7 @@ def real_transform(_):
             return new_updates.delta_updates, new_state
 
         def stop_transform(_):
-            return jax.tree_map(jnp.zeros_like, updates.delta_updates), prev_state
+            return jax.tree_util.tree_map(jnp.zeros_like, updates.delta_updates), prev_state
 
         # We do the computation regardless of the should_update value, so we could have
         # equally used jnp.where() here instead.

axlearn/common/utils.py

@@ -1991,7 +1991,7 @@ def validate_contains_paths(x: Nested[Tensor], paths: Sequence[str]):
         except KeyError as e:
             raise ValueError(
                 f"Input is expected to contain '{path}'; "
-                f"instead, it contains: '{jax.tree_structure(x)}'."
+                f"instead, it contains: '{jax.tree_util.tree_structure(x)}'."
             ) from e
 
 

axlearn/experiments/text/gpt/common_test.py

@@ -52,7 +52,9 @@ def test_mesh_axes(self):
         # axis for multiple dims.
         for v in cfg.input.input_partitioner.path_rank_to_partition.values():
             visited = set()
-            for axis in jax.tree_leaves(tuple(v)):  # Cast to tuple since PartitionSpec is a leaf.
+            for axis in jax.tree_util.tree_leaves(
+                tuple(v)
+            ):  # Cast to tuple since PartitionSpec is a leaf.
                 self.assertNotIn(axis, visited)
                 visited.add(axis)
             self.assertGreater(len(visited), 0)