Функция take

Функция take возвращает новый тензор, содержащий элементы исходного тензора, расположенные по указанным индексам. Индексы задаются в виде одномерного тензора и интерпретируются как плоские индексы, то есть элементы исходного тензора рассматриваются как одномерный массив. Первым параметром функция принимает исходный тензор, вторым - тензор с индексами выбираемых элементов.

Синтаксис

torch.take(input, indices)

Пример

Давайте создадим тензор из чисел и выберем из него элементы по индексам 0, 2, 4:

import torch

t = torch.tensor([10, 20, 30, 40, 50])
indices = torch.tensor([0, 2, 4])
res = torch.take(t, indices)
print(res)

Результат выполнения кода:

tensor([10, 30, 50])

Пример

Теперь давайте создадим двумерный тензор и выберем из него элементы по плоским индексам 1, 3, 5. Обратите внимание, что индексы считаются от нуля по всем элементам, развернутым в одномерный массив:

import torch

t = torch.tensor([
	[1, 2, 3],
	[4, 5, 6],
])
indices = torch.tensor([1, 3, 5])
res = torch.take(t, indices)
print(res)

Результат выполнения кода:

tensor([2, 4, 6])

Пример

Функция take также поддерживает использование отрицательных индексов, которые отсчитываются от конца плоского представления тензора:

import torch

t = torch.tensor([10, 20, 30, 40, 50])
indices = torch.tensor([-1, -3, -5])
res = torch.take(t, indices)
print(res)

Результат выполнения кода:

tensor([50, 30, 10])

Пример

Если индекс выходит за пределы допустимого диапазона, функция выбрасывает ошибку. Давайте попробуем выбрать элемент с индексом, превышающим количество элементов в тензоре:

import torch

t = torch.tensor([1, 2, 3])
indices = torch.tensor([5])
res = torch.take(t, indices)

Результат выполнения кода:

"IndexError: index 5 is out of bounds for dimension 0 with size 3"

Смотрите также

функцию index_select,
которая выбирает элементы тензора по индексам вдоль указанной размерности
функцию gather,
которая собирает элементы тензора по индексам из другой размерности
функцию nonzero,
которая возвращает индексы ненулевых элементов тензора
функцию where,
которая выбирает элементы тензора по условию