maxframe.tensor.flatnonzero

maxframe.tensor.flatnonzero(a)

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

这等价于 a.ravel().nonzero()[0]。

参数:

a (Tensor) -- 输入张量。

返回:

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

返回类型:

Tensor

参见

nonzero

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

ravel

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

示例

>>> import maxframe.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()