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

google-research / scenic

Scenic: A Jax Library for Computer Vision Research and Beyond

Apache License 2.0

3.14k stars 417 forks source link

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

Closed copybara-service[bot] closed 3 months ago

copybara-service[bot] commented 3 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).