jax.random.permutation doesn't fully replace jax.random.shuffle

[maurorigo]

In a code that I'm writing, I need to permute a matrix along the rows independently from one column to another. jax.random.shuffle allows to do this in one application, but it is deprecated and the warning message suggests using jax.random.permutation, which, at least to my knowledge, does not do that. For example, applying jax.random.shuffle to the matrix [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]] yields for instance [[3 2 3 2], [2 3 1 3], [1 1 2 1]], while jax.random.permutation gives for instance [[3 3 3 3], [2 2 2 2], [1 1 1 1]]. I could still use jax.random.permutation and loop over the columns, but it is less efficient. Therefore I would like to kindly ask you if there is another way to do this permutation efficiently without using jax.random.shuffle.

[jakevdp]

I think the reason jax.random.shuffle is deprecated is because it is potentially confusing: numpy.random.shuffle permutes contents in-place, and does not permute independently along each row, so having a function of the same name in jax.random with markedly different behavior might be confusing.

For what it's worth, you can recover the independent shuffling behavior by using permutation with vmap; for example:

import jax
import jax.numpy as jnp
from jax import random

key = random.PRNGKey(42)
x = jnp.broadcast_to(jnp.arange(3), (4, 3))

print(random.shuffle(key, x, axis=1))
# [[0 2 1]
#  [1 0 2]
#  [1 2 0]
#  [2 1 0]]
print(jax.vmap(random.permutation)(random.split(key, x.shape[0]), x))
# [[1 2 0]
#  [0 2 1]
#  [2 1 0]
#  [1 0 2]]

... though it's admittedly less convenient.

Perhaps the current shuffle could be replaced by adding an independent keyword to permutation that defaults to False? What do you think?

[maurorigo]

Oh that's true, I didn't think of using vmap. The solution of implementing an independent keyword would be perfect in my opinion! Thank you very much for your answer.

[jakevdp]

Great - keyword added in #8430