Replace deprecated jax.random.shuffle with jax.random.permutation

jax.random.shuffle has long been deprecated, because it cannot operate in-place like np.random.shuffle, and because its functionality can be performed with jax.random.permutation (with independent=True in the case of multi-dimensional arrays). PiperOrigin-RevId: 622218990
google-research · Apr 5, 2024 · 78e8a4f · 78e8a4f
1 parent e895bbf
commit 78e8a4f
Show file tree

Hide file tree

Showing 27 changed files with 27 additions and 29 deletions.
diff --git a/precondition/__init__.py b/precondition/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/distributed_shampoo.py b/precondition/distributed_shampoo.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/distributed_shampoo_test.py b/precondition/distributed_shampoo_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/oco/algorithms.py b/precondition/oco/algorithms.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/oco/datasets.py b/precondition/oco/datasets.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/oco/sweep.py b/precondition/oco/sweep.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/oco/train.py b/precondition/oco/train.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/quantization_utils.py b/precondition/quantization_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/sm3.py b/precondition/sm3.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/sm3_test.py b/precondition/sm3_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/grafting.py b/precondition/tearfree/grafting.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/grafting_test.py b/precondition/tearfree/grafting_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/momentum.py b/precondition/tearfree/momentum.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/momentum_test.py b/precondition/tearfree/momentum_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/optimizer.py b/precondition/tearfree/optimizer.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/optimizer_smoke_test.py b/precondition/tearfree/optimizer_smoke_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/optimizer_test.py b/precondition/tearfree/optimizer_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/praxis_shim.py b/precondition/tearfree/praxis_shim.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/reallocation.py b/precondition/tearfree/reallocation.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/reallocation_test.py b/precondition/tearfree/reallocation_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/reshaper.py b/precondition/tearfree/reshaper.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/reshaper_test.py b/precondition/tearfree/reshaper_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/second_order.py b/precondition/tearfree/second_order.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/shampoo.py b/precondition/tearfree/shampoo.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition 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/precondition/tearfree/shampoo_test.py b/precondition/tearfree/shampoo_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The precondition Authors.
+# Copyright 2024 The precondition Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -289,8 +289,6 @@ def test_precond_freq(self):
     rng1, rng2 = jax.random.split(key, 2)
     seq1 = jnp.arange(freq, dtype=int)
     seq2 = jnp.copy(seq1)
-    jax.random.shuffle(rng1, seq1)
-    jax.random.shuffle(rng2, seq2)
     # Shuffle within groups of <freq>
     shuffled = jnp.take(grads, jnp.concatenate([seq1, seq2 + freq]), axis=0)
 

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