jax.random.permutation

Contents

jax.random.permutation#

jax.random.permutation(key, x, axis=0, independent=False)[source]#

Returns a randomly permuted array or range.

Parameters:
  • key (KeyArrayLike) – 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.

Return type:

Array

Returns:

A shuffled version of x or array range