mars.tensor.argwhere

mars.tensor.argwhere(a)

Find the indices of tensor elements that are non-zero, grouped by element.

参数

a (array_like) – Input data.

返回

index_tensor – Indices of elements that are non-zero. Indices are grouped by element.

返回类型

Tensor

参见

where, nonzero

备注

mt.argwhere(a) is the same as mt.transpose(mt.nonzero(a)).

The output of argwhere is not suitable for indexing tensors. For this purpose use nonzero(a) instead.

示例

>>> import mars.tensor as mt

>>> x = mt.arange(6).reshape(2,3)
>>> x.execute()
array([[0, 1, 2],
       [3, 4, 5]])
>>> mt.argwhere(x>1).execute()
array([[0, 2],
       [1, 0],
       [1, 1],
       [1, 2]])