oneflow.argwhere

oneflow.argwhere(input, dtype: Optional[oneflow._oneflow_internal.dtype] = oneflow.int32)

This operator finds the indices of input Tensor input elements that are non-zero.

It returns a list in which each element is a coordinate that points to a non-zero element in the condition.

Parameters
  • input (oneflow.Tensor) – the input Tensor.

  • dtype (Optional[flow.dtype], optional) – The data type of output. Defaults to None.

Returns

The result Tensor.

Return type

oneflow.Tensor

For example:

>>> import numpy as np
>>> import oneflow as flow
>>> x = np.array([[0, 1, 0],
...            [2, 0, 2]]).astype(np.float32)

>>> input = flow.Tensor(x)
>>> output = flow.argwhere(input)
>>> output
tensor([[0, 1],
        [1, 0],
        [1, 2]], dtype=oneflow.int32)