torch.nonzero

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

Примечание

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

torch.nonzero(..., as_tuple=True) возвращает кортеж 1-мерных тензоров индексов, позволяющий выполнять продвинутое индексирование, поэтому 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:

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

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

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

Параметры:

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

Ключевые аргументы:

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

Возвращает:

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

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

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]),)