torch.nonzero (original) (raw)

Note

torch.nonzero(..., as_tuple=False) (default) returns a 2-D tensor where each row is the index for a nonzero value.

torch.nonzero(..., as_tuple=True) returns a tuple of 1-D index tensors, allowing for advanced indexing, so x[x.nonzero(as_tuple=True)]gives all nonzero values of tensor x. Of the returned tuple, each index tensor contains nonzero indices for a certain dimension.

See below for more details on the two behaviors.

When input is on CUDA, torch.nonzero() causes host-device synchronization.

When as_tuple is False (default):

Returns a tensor containing the indices of all non-zero elements ofinput. Each row in the result contains the indices of a non-zero element in input. The result is sorted lexicographically, with the last index changing the fastest (C-style).

If input has nn dimensions, then the resulting indices tensorout is of size (z×n)(z \times n), where zz is the total number of non-zero elements in the input tensor.

When as_tuple is True:

Returns a tuple of 1-D tensors, one for each dimension in input, each containing the indices (in that dimension) of all non-zero elements ofinput .

If input has nn dimensions, then the resulting tuple contains nntensors of size zz, where zz is the total number of non-zero elements in the input tensor.

As a special case, when input has zero dimensions and a nonzero scalar value, it is treated as a one-dimensional tensor with one element.

Parameters

input (Tensor) – the input tensor.

Keyword Arguments

out (LongTensor , optional) – the output tensor containing indices

Returns

If as_tuple is False, the output tensor containing indices. If as_tuple is True, one 1-D tensor for each dimension, containing the indices of each nonzero element along that dimension.

Return type

LongTensor or tuple of LongTensor

Example:

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