mars.tensor.flatnonzero

mars.tensor.flatnonzero(a)

返回在展平版本的 a 中非零的索引。

这相当于 a.ravel().nonzero()[0]。

Parameters

a (Tensor) – 输入张量。

Returns

res – 输出张量，包含了a.ravel()中非零元素的索引。

Return type

张量

另请参阅

nonzero

返回输入张量中非零元素的索引。

ravel

返回一个包含输入张量元素的一维张量。

示例

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

使用非零元素的索引作为索引数组来提取这些元素：

>>> x.ravel()[mt.flatnonzero(x)].execute()  # TODO(jisheng): accomplish this after fancy indexing is supported