shuffle

Contents

shuffle#

class brainstate.random.shuffle(x, axis=0, key=None)#

Return a randomly shuffled copy of an array along an axis.

The contents of x are permuted along the given axis (only the order of slices along that axis changes; the values within each slice are kept).

Note

Unlike numpy.random.shuffle(), this function does not modify x in place. JAX arrays are immutable, so a new array is returned and the input is left unchanged. Bind the result to use it.

Parameters:
  • x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – The array (or array-like) to shuffle. May carry physical units.

  • axis (int) – The axis along which x is shuffled.

  • key (int | Array | ndarray | None) – Seed or PRNG key for the random number generator. If not given, the default random number generator is used.

Returns:

out – A new array with the same shape and dtype as x, shuffled along axis. A Quantity is returned when x carries physical units.

Return type:

Array | Quantity

Examples

Shuffle a 1D array (the input is not modified):

>>> import brainstate
>>> import numpy as np
>>> arr = np.arange(10)
>>> shuffled = brainstate.random.shuffle(arr)
>>> print(set(np.asarray(shuffled)) == set(arr))  # True (same elements)

Multi-dimensional arrays are shuffled along axis (default first axis):

>>> arr = np.arange(9).reshape((3, 3))
>>> shuffled = brainstate.random.shuffle(arr)
>>> print(shuffled.shape == arr.shape)  # True (shape preserved)
>>> print(sorted(np.asarray(shuffled).flatten()) == list(range(9)))  # True