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: 635648241
magenta · May 21, 2024 · 7996dc5 · 7996dc5
1 parent a8bab04
commit 7996dc5
Show file tree

Hide file tree

Showing 30 changed files with 33 additions and 33 deletions.
diff --git a/music_spectrogram_diffusion/__init__.py b/music_spectrogram_diffusion/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/audio_codecs.py b/music_spectrogram_diffusion/audio_codecs.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/beam/evaluation.py b/music_spectrogram_diffusion/beam/evaluation.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -206,7 +206,7 @@ def process(self, id_examples):
 
       tick = time.time()
       if self._prediction_source == 'model':
-        pred_encoded, _ = jax.tree_map(
+        pred_encoded, _ = jax.tree.map(
             np.asarray, self._synth_model.predict(batch))
       elif self._prediction_source in ('gt_raw', 'gt_encoded'):
         pred_encoded = batch['decoder_target_tokens']

diff --git a/music_spectrogram_diffusion/datasets.py b/music_spectrogram_diffusion/datasets.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/dump_task.py b/music_spectrogram_diffusion/dump_task.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/event_codec.py b/music_spectrogram_diffusion/event_codec.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/event_codec_test.py b/music_spectrogram_diffusion/event_codec_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/feature_converters.py b/music_spectrogram_diffusion/feature_converters.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/inference.py b/music_spectrogram_diffusion/inference.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/layers.py b/music_spectrogram_diffusion/layers.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/layers_test.py b/music_spectrogram_diffusion/layers_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/metrics.py b/music_spectrogram_diffusion/metrics.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/metrics_test.py b/music_spectrogram_diffusion/metrics_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/models/autoregressive/models.py b/music_spectrogram_diffusion/models/autoregressive/models.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/models/autoregressive/network.py b/music_spectrogram_diffusion/models/autoregressive/network.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/models/autoregressive/output_functions.py b/music_spectrogram_diffusion/models/autoregressive/output_functions.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/models/diffusion/diffusion_utils.py b/music_spectrogram_diffusion/models/diffusion/diffusion_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/models/diffusion/feature_converters.py b/music_spectrogram_diffusion/models/diffusion/feature_converters.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/models/diffusion/models.py b/music_spectrogram_diffusion/models/diffusion/models.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -178,7 +178,7 @@ def predict_batch_with_aux(
     def pred_fn(z: jnp.ndarray,
                 time: jnp.ndarray,
                 include_conditioning: bool) -> jnp.ndarray:
-      step_encodings_and_masks = jax.tree_map(
+      step_encodings_and_masks = jax.tree.map(
           lambda x: x * include_conditioning, encodings_and_masks)
       return self.module.apply(
           {
@@ -373,7 +373,7 @@ def predict_batch_with_aux(
     def pred_fn(z: jnp.ndarray,
                 time: jnp.ndarray,
                 include_conditioning: bool) -> jnp.ndarray:
-      step_encodings_and_masks = jax.tree_map(
+      step_encodings_and_masks = jax.tree.map(
           lambda x: x * include_conditioning, encodings_and_masks)
       return self.module.apply(
           {

diff --git a/music_spectrogram_diffusion/models/diffusion/network.py b/music_spectrogram_diffusion/models/diffusion/network.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/note_sequences.py b/music_spectrogram_diffusion/note_sequences.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/note_sequences_test.py b/music_spectrogram_diffusion/note_sequences_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/postprocessors.py b/music_spectrogram_diffusion/postprocessors.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/preprocessors.py b/music_spectrogram_diffusion/preprocessors.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/run_length_encoding.py b/music_spectrogram_diffusion/run_length_encoding.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/run_length_encoding_test.py b/music_spectrogram_diffusion/run_length_encoding_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/tasks.py b/music_spectrogram_diffusion/tasks.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/transcription_inference.py b/music_spectrogram_diffusion/transcription_inference.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/vocabularies.py b/music_spectrogram_diffusion/vocabularies.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion 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/music_spectrogram_diffusion/vocabularies_test.py b/music_spectrogram_diffusion/vocabularies_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The Music Spectrogram Diffusion Authors.
+# Copyright 2024 The Music Spectrogram Diffusion Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.