torch.unravel_index

torch.unravel_index(indices, shape) [исходный код]

Преобразует тензор плоских индексов в кортеж тензоров координат, индексирующих произвольный тензор заданной формы.

Параметры:

• indices (Tensor) – Целочисленный тензор, содержащий индексы в уплощённой версии произвольного тензора формы shape. Все элементы должны находиться в диапазоне [0, prod(shape) - 1].
• shape (int, sequence of ints, or torch.Size) – Форма произвольного тензора. Все элементы должны быть неотрицательными.

Возвращает:

Каждый i-й тензор в результате соответствует измерению i тензора shape. Каждый тензор имеет ту же форму, что и indices, и содержит по одному индексу измерения i для каждого плоского индекса, заданного в indices.

Тип возвращаемого значения:

tuple тензоров

Пример:

>>> import torch
>>> torch.unravel_index(torch.tensor(4), (3, 2))
(tensor(2),
 tensor(0))

>>> torch.unravel_index(torch.tensor([4, 1]), (3, 2))
(tensor([2, 0]),
 tensor([0, 1]))

>>> torch.unravel_index(torch.tensor([0, 1, 2, 3, 4, 5]), (3, 2))
(tensor([0, 0, 1, 1, 2, 2]),
 tensor([0, 1, 0, 1, 0, 1]))

>>> torch.unravel_index(torch.tensor([1234, 5678]), (10, 10, 10, 10))
(tensor([1, 5]),
 tensor([2, 6]),
 tensor([3, 7]),
 tensor([4, 8]))

>>> torch.unravel_index(torch.tensor([[1234], [5678]]), (10, 10, 10, 10))
(tensor([[1], [5]]),
 tensor([[2], [6]]),
 tensor([[3], [7]]),
 tensor([[4], [8]]))

>>> torch.unravel_index(torch.tensor([[1234], [5678]]), (100, 100))
(tensor([[12], [56]]),
 tensor([[34], [78]]))