Выбор по номерам в TensorFlow

Чтобы взять элементы не подряд, а по списку номеров, служит функция gather. Она читает исходный тензор и возвращает новый в том порядке, в каком заданы номера. Один номер можно указать несколько раз.

Создадим ряд из четырёх чисел и соберём элементы в порядке 3, 0 и 2:

import tensorflow as tf

row = tf.constant([5, 10, 15, 20])
picked = tf.gather(row, [3, 0, 2])
print(picked)
# выведет tf.Tensor([20  5 15], shape=(3,), dtype=int32)

В таблице номера задают, какие строки или столбцы взять целиком. Параметр axis указывает, вдоль какой оси идёт выбор. Зададим таблицу 3×3 и возьмём строки 2 и 0:

import tensorflow as tf

table = tf.constant([
	[1, 2, 3],
	[4, 5, 6],
	[7, 8, 9],
])
rows = tf.gather(table, [2, 0], axis=0)
print(rows)
# выведет tf.Tensor([[7 8 9] [1 2 3]], shape=(2, 3), dtype=int32)

Если указать axis равным 1, из каждой строки берут выбранные столбцы. Возьмём столбцы 2 и 0:

import tensorflow as tf

table = tf.constant([
	[1, 2, 3],
	[4, 5, 6],
	[7, 8, 9],
])
cols = tf.gather(table, [2, 0], axis=1)
print(cols)
# выведет tf.Tensor([[3 1] [6 4] [9 7]], shape=(3, 2), dtype=int32)

Из ряда [100, 200, 300, 400, 500] соберите новый ряд по номерам [4, 1, 1, 0] и выведите его.

Из ряда [7, 8, 9] получите ряд по номерам [2, 2, 0] и выведите длину результата.

Соберите таблицу из трёх строк [1, 2, 3], [4, 5, 6] и [7, 8, 9], возьмите из неё первую и третью строки и выведите таблицу из двух строк.