oneflow.nonzero

oneflow.nonzero(input, *, out=None, as_tuple=False)Tensor or tuple of Tensors

Note

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

When as_tuple is 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 as_tuple is False (default):

Returns a tensor containing the indices of all non-zero elements of input. 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 \(n\) dimensions, then the resulting indices tensor out is of size \((z \times n)\), where \(z\) 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 of input .

If input has \(n\) dimensions, then the resulting tuple contains \(n\) tensors of size \(z\), where \(z\) 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 (Tensor, 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

Tensor or tuple of Tensors

Example:

>>> import oneflow as flow
>>> flow.nonzero(flow.tensor([1, 1, 1, 0, 1]))
tensor([[0],
        [1],
        [2],
        [4]], dtype=oneflow.int64)
>>> flow.nonzero(flow.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]], dtype=oneflow.int64)
>>> flow.nonzero(flow.tensor([1, 1, 1, 0, 1]), as_tuple=True)
(tensor([0, 1, 2, 4], dtype=oneflow.int64),)
>>> flow.nonzero(flow.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], dtype=oneflow.int64), tensor([0, 1, 2, 3], dtype=oneflow.int64))
>>> flow.nonzero(flow.tensor(5), as_tuple=True)
(tensor([0], dtype=oneflow.int64),)