torch.nonzero

torch.nonzero(input, *, out=None, as_tuple=False) → LongTensor or tuple of LongTensors

Примечание

torch.nonzero(..., as_tuple=False) (по умолчанию) возвращает двумерный тензор, в каждой строке которого находится индекс ненулевого значения.

torch.nonzero(..., as_tuple=True) возвращает кортеж одномерных тензоров индексов, что позволяет использовать расширенную индексацию, поэтому x[x.nonzero(as_tuple=True)] возвращает все ненулевые значения тензора x. Каждый тензор индексов в возвращаемом кортеже содержит ненулевые индексы для определённого измерения.

Подробнее об этих двух режимах см. ниже.

Если input находится на CUDA, torch.nonzero() вызывает синхронизацию между хостом и устройством.

Если as_tuple имеет значение False (по умолчанию):

Возвращает тензор, содержащий индексы всех ненулевых элементов input. Каждая строка результата содержит индексы ненулевого элемента input. Результат отсортирован лексикографически: последний индекс изменяется быстрее всего (в стиле C).

Если input имеет $n$ измерений, то результирующий тензор индексов out имеет размер $(z \times n)$, где $z$ — общее количество ненулевых элементов в тензоре input.

Если as_tuple имеет значение True:

Возвращает кортеж одномерных тензоров — по одному для каждого измерения input; каждый содержит индексы (в этом измерении) всех ненулевых элементов input .

Если input имеет $n$ измерений, результирующий кортеж содержит $n$ тензоров размером $z$, где $z$ — общее количество ненулевых элементов в тензоре input.

В частном случае, когда input имеет нулевое число измерений и ненулевое скалярное значение, он рассматривается как одномерный тензор с одним элементом.

Параметры:

input (Tensor) – входной тензор.

Именованные аргументы:

out (LongTensor, необязательно) – выходной тензор, содержащий индексы

Возвращает:

Если as_tuple имеет значение False, выходной тензор, содержащий индексы. Если as_tuple имеет значение True, по одному одномерному тензору для каждого измерения, содержащему индексы каждого ненулевого элемента вдоль этого измерения.

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

LongTensor или кортеж из LongTensor

Пример:

>>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1]))
tensor([[ 0],
        [ 1],
        [ 2],
        [ 4]])
>>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0],
...                             [0.0, 0.4, 0.0, 0.0],
...                             [0.0, 0.0, 1.2, 0.0],
...                             [0.0, 0.0, 0.0,-0.4]]))
tensor([[ 0,  0],
        [ 1,  1],
        [ 2,  2],
        [ 3,  3]])
>>> torch.nonzero(torch.tensor([1, 1, 1, 0, 1]), as_tuple=True)
(tensor([0, 1, 2, 4]),)
>>> torch.nonzero(torch.tensor([[0.6, 0.0, 0.0, 0.0],
...                             [0.0, 0.4, 0.0, 0.0],
...                             [0.0, 0.0, 1.2, 0.0],
...                             [0.0, 0.0, 0.0,-0.4]]), as_tuple=True)
(tensor([0, 1, 2, 3]), tensor([0, 1, 2, 3]))
>>> torch.nonzero(torch.tensor(5), as_tuple=True)
(tensor([0]),)