jax.random.permutation#
- jax.random.permutation(key, x, axis=0, independent=False, *, out_sharding=None)[source]#
Returns a randomly permuted array or range.
- Parameters:
key (ArrayLike) – a PRNG key used as the random key.
x (int | ArrayLike) – int or array. If x is an integer, randomly shuffle np.arange(x). If x is an array, randomly shuffle its elements.
axis (int) – int, optional. The axis which x is shuffled along. Default is 0.
independent (bool) – bool, optional. If set to True, each individual vector along the given axis is shuffled independently. Default is False.
- Returns:
A shuffled version of x or array range
- Return type: