mars.tensor.
flatnonzero
Return indices that are non-zero in the flattened version of a.
This is equivalent to a.ravel().nonzero()[0].
a (Tensor) – Input tensor.
res – Output tensor, containing the indices of the elements of a.ravel() that are non-zero.
Tensor
See also
nonzero
Return the indices of the non-zero elements of the input tensor.
ravel
Return a 1-D tensor containing the elements of the input tensor.
Examples
>>> import mars.tensor as mt
>>> x = mt.arange(-2, 3) >>> x.execute() array([-2, -1, 0, 1, 2]) >>> mt.flatnonzero(x).execute() array([0, 1, 3, 4])
Use the indices of the non-zero elements as an index array to extract these elements:
>>> x.ravel()[mt.flatnonzero(x)].execute() # TODO(jisheng): accomplish this after fancy indexing is supported