megengine.functional.non_zero¶
- non_zero(condition, as_tuple=False)[source]¶
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- Returns
one tuple of 1-D tensors or one tensor
Examples
>>> 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),)