tf.random.experimental.index_shuffle

Возвращает позицию index в перестановке [0, ..., max_index].

tf.random.experimental.index_shuffle(
    index, seed, max_index
)

Для каждой возможной пары seed и max_index существует одна псевдослучайная перестановка последовательности S=[0, ..., max_index]. Вместо материализации всего массива мы можем вычислить новую позицию любого целого числа i (0 <= i <= max_index) в S. Это может быть полезно для очень больших max_index, поскольку позволяет избежать выделения больших объемов памяти.

В простейшем случае index и max_index — скаляры, а seed — вектор длины 2 (как обычно бывает для RNG без состояния). Однако ко всем ним можно добавить ведущую размерность пакета. Если у некоторых из них размерности пакета нет, а у других есть, index_shuffle добавит размерность пакета к первым с помощью широковещательного распространения.

Входное значение index и выходные данные можно использовать в качестве индексов для перемешивания вектора. Например:

vector = tf.constant(['e0', 'e1', 'e2', 'e3'])
indices = tf.random.experimental.index_shuffle(
  index=tf.range(4), seed=[5, 9], max_index=3)
print(indices)
tf.Tensor([2 0 1 3], shape=(4,), dtype=int32)
shuffled_vector = tf.gather(vector, indices)
print(shuffled_vector)
tf.Tensor([b'e2' b'e0' b'e1' b'e3'], shape=(4,), dtype=string)

Что еще полезнее, это можно использовать в потоковых (то есть онлайн-) сценариях, например, с tf.data, где каждый элемент vector обрабатывается отдельно, а весь vector никогда не материализуется в памяти.

dataset = tf.data.Dataset.range(10)
dataset = dataset.map(
 lambda idx: tf.random.experimental.index_shuffle(idx, [5, 8], 9))
print(list(dataset.as_numpy_iterator()))
[3, 8, 0, 1, 2, 7, 6, 9, 4, 5]

Эта операция не имеет состояния (как функции tf.random.stateless_*), то есть выходные данные полностью определяются значением seed (при прочих равных условиях). Каждому выбору seed соответствует одна перестановка, поэтому при многократном вызове этой функции для одного и того же перемешивания используйте одно и то же значение seed. Например:

seed = [5, 9]
idx0 = tf.random.experimental.index_shuffle(0, seed, 3)
idx1 = tf.random.experimental.index_shuffle(1, seed, 3)
idx2 = tf.random.experimental.index_shuffle(2, seed, 3)
idx3 = tf.random.experimental.index_shuffle(3, seed, 3)
shuffled_vector = tf.gather(vector, [idx0, idx1, idx2, idx3])
print(shuffled_vector)
tf.Tensor([b'e2' b'e0' b'e1' b'e3'], shape=(4,), dtype=string)