search

google / uncertainty-baselines

High-quality implementations of standard and SOTA methods on a variety of tasks.
Apache License 2.0
1.43k stars 202 forks source link

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

Closed copybara-service[bot] closed 5 months ago

copybara-service[bot] commented 5 months ago

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).