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 时

返回一个由 1-D 张量组成的元组，每个张量对应 input 中的一个维度，其中包含 input 中所有非零元素（在该维度上）的索引。

如果 input 有 $n$ 个维度，则返回的元组包含 $n$ 个大小为 $z$ 的张量，其中 $z$ 是 input 张量中非零元素的总数。

作为特例，当 input 的维度为零且其值为非零标量时，它被视为一个只有一个元素的 1 维张量。

参数

input (Tensor) – 输入张量。

关键字参数

out (LongTensor, optional) – 包含索引的输出张量

返回

如果 as_tuple 为 False，则为包含索引的输出张量。如果 as_tuple 为 True，则为每个维度创建一个 1-D 张量，包含沿该维度每个非零元素的索引。

返回类型

LongTensor 或 LongTensor 的 tuple

示例

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