Skip to content

← JAX Stochasticity step 23 of 25

Easy Primitives

Random Permutation

Jump to the task

Reorder values without losing any

A permutation rearranges positions. Every input occurrence appears exactly once in the output, including repeated values. Shuffling [1, 1, 2] still gives two ones and one two. It does not sample with replacement, remove duplicates or generate new values.

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
x = jnp.array([1., 2., 3., 4., 5.])
shuffled = jax.random.permutation(key, x)
# [1., 2., 5., 4., 3.] with the exercise runtime's default settings
# x remains [1., 2., 3., 4., 5.]

Some positions can stay where they were. Even the unchanged order is a possible random permutation. Different keys request different draws, but do not guarantee different visible outputs, especially for small arrays or arrays with repeated values.

Keep related data aligned

For training data, generate one permutation of row indices and apply it to both features and labels. For example, order [2, 0, 1] sends features [10, 20, 30] to [30, 10, 20] and labels [0, 1, 2] to [2, 0, 1]. Shuffling them independently can pair each feature with the wrong label.

order = jax.random.permutation(epoch_key, features.shape[0])
shuffled_features = features[order]
shuffled_labels = labels[order]

This snippet assumes epoch_key, features and labels are already defined. Derive a fresh epoch key with split or fold_in when a new order is wanted. A fixed key and sampling configuration let you reproduce an order within the same JAX setup; matching integer seeds across frameworks does not guarantee matching permutations.

Multidimensional arrays need an axis

The default axis is 0. With independent=False, a matrix's rows move together as whole slices. With axis=1, columns move together instead. With independent=True, each slice along the chosen axis is shuffled separately. To shuffle within each row independently, specify both axis=1 and independent=True. independent=True alone keeps axis=0 and shuffles down each column separately. This exercise uses a vector.

Problem

Implement shuffle_array(seed, x) by creating jax.random.PRNGKey(int(seed)) and returning jax.random.permutation(key, x).

Seed is a whole Python number in [0, 2**32), supplied as a float. x is a one-dimensional numeric JAX array. Preserve its shape, dtype and values, including duplicates. Empty and one-element inputs keep their values. Return the result; the operation does not mutate x.

Use the default runtime Threefry configuration and the original seed-derived key. Tests check its reproducible order, not just whether the result contains the input values.

For JIT, mark seed static because this interface uses int(seed): jax.jit(shuffle_array, static_argnames=("seed",)). The input array can remain dynamic.