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 16, 2024 · 3819760 · 3819760
1 parent 560ad94
commit 3819760
Show file tree

Hide file tree

Showing 29 changed files with 30 additions and 30 deletions.
diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/__init__.py b/mt3/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/datasets.py b/mt3/datasets.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/event_codec.py b/mt3/event_codec.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/event_codec_test.py b/mt3/event_codec_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/inference.py b/mt3/inference.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/layers.py b/mt3/layers.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/layers_test.py b/mt3/layers_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -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': [[

diff --git a/mt3/metrics.py b/mt3/metrics.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/metrics_utils.py b/mt3/metrics_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/metrics_utils_test.py b/mt3/metrics_utils_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/mixing.py b/mt3/mixing.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/models.py b/mt3/models.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/network.py b/mt3/network.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/note_sequences.py b/mt3/note_sequences.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/note_sequences_test.py b/mt3/note_sequences_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/preprocessors.py b/mt3/preprocessors.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/run_length_encoding.py b/mt3/run_length_encoding.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/run_length_encoding_test.py b/mt3/run_length_encoding_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/scripts/dump_task.py b/mt3/scripts/dump_task.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/scripts/extract_monophonic_examples.py b/mt3/scripts/extract_monophonic_examples.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/spectral_ops.py b/mt3/spectral_ops.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/spectrograms.py b/mt3/spectrograms.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/summaries.py b/mt3/summaries.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/tasks.py b/mt3/tasks.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/version.py b/mt3/version.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/vocabularies.py b/mt3/vocabularies.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/mt3/vocabularies_test.py b/mt3/vocabularies_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/setup.py b/setup.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The MT3 Authors.
+# Copyright 2024 The MT3 Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.