megengine.functional.non_zero

non_zero(condition, as_tuple=False)[源代码]

When as_tuple is False (default): Returns a tensor including the indices of all non-zero elements of Tensor condition. Every row in the result including the indices of a non-zero element in input. The result is sorted in lexicography order, with the last index changing the fastest (C-style). When as_tuple is True: Returns a tuple of 1-D tensors, one for each dimension in input, each containing the indices (in that dimension) of all non-zero elements of condition. :type condition: Tensor :param condition: :type condition: Tensor

返回

one tuple of 1-D tensors or one tensor

实际案例

>>> import numpy as np
>>> condition = Tensor(np.array([1,1,0,1]))
>>> index = F.non_zero(condition,as_tuple=True)
>>> print(index)
(Tensor([0 1 3], dtype=int32, device=xpux:0),)