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: 633777708
magenta · May 22, 2024 · 3c6b607 · 3c6b607
1 parent 2f8e2c8
commit 3c6b607
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mt3/layers_test.py b/mt3/layers_test.py
@@ -500,7 +500,7 @@ def test_mlp_same_out_dim(self):
         dtype=np.float32)
     params = module.init(random.PRNGKey(0), inputs, deterministic=True)
     self.assertEqual(
-        jax.tree_map(lambda a: a.tolist(), params), {
+        jax.tree.map(lambda a: a.tolist(), params), {
             'params': {
                 'wi': {
                     'kernel': [[