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

Функция index_select собирает срез вдоль одной оси: в параметре dim указывают, по рядам или по столбцам идёт выбор, а во втором тензоре перечисляют номера, которые нужно взять.

Создадим таблицу из трёх рядов и оставим только первый и третий:

import torch

scores = torch.tensor([
	[10, 11, 12],
	[20, 21, 22],
	[30, 31, 32],
])
row_ids = torch.tensor([0, 2])
picked = torch.index_select(scores, 0, row_ids)
print(picked)

Тот же приём работает по столбцам: номера задают, какие столбцы попадут в результат и в каком порядке:

import torch

scores = torch.tensor([
	[10, 11, 12],
	[20, 21, 22],
])
col_ids = torch.tensor([2, 0])
reordered = torch.index_select(scores, 1, col_ids)
print(reordered) # выведет tensor([[12, 10], [22, 20]])

Создайте таблицу из рядов [1, 2], [3, 4] и [5, 6]; по номерам 0 и 2 соберите два ряда и выведите результат.

Создайте таблицу 2 на 3 из чисел от 1 до 6 по рядам; возьмите столбцы с номерами 2 и 0 и выведите таблицу шириной 2.

Создайте таблицу из четырёх рядов по одному числу: 10, 20, 30, 40; оставьте ряды с номерами 1 и 3 и выведите число рядов в ответе.