r"""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
(condition: Tensor, as_tuple=False)
| 932 | |
| 933 | |
| 934 | def non_zero(condition: Tensor, as_tuple=False): |
| 935 | r"""When as_tuple is False (default): |
| 936 | Returns a tensor including the indices of all non-zero elements of Tensor condition. |
| 937 | Every row in the result including the indices of a non-zero element in input. |
| 938 | The result is sorted in lexicography order, with the last index changing the fastest (C-style). |
| 939 | When as_tuple is True: |
| 940 | Returns a tuple of 1-D tensors, one for each dimension in input, |
| 941 | each containing the indices (in that dimension) of all non-zero elements of condition. |
| 942 | Args: |
| 943 | condition(Tensor) - the input tensor |
| 944 | Returns: |
| 945 | one tuple of 1-D tensors or one tensor |
| 946 | |
| 947 | Examples: |
| 948 | >>> import numpy as np |
| 949 | >>> condition = Tensor(np.array([1,1,0,1])) |
| 950 | >>> index = F.non_zero(condition,as_tuple=True) |
| 951 | >>> print(index) |
| 952 | (Tensor([0 1 3], dtype=int32, device=xpux:0),) |
| 953 | """ |
| 954 | |
| 955 | if not isinstance(condition, Tensor): |
| 956 | raise TypeError("input must be a tensor") |
| 957 | op = builtin.NonZero() |
| 958 | (index,) = apply(op, condition) |
| 959 | ret = None |
| 960 | if as_tuple == True: |
| 961 | arr = [] |
| 962 | for index_ele in range(0, condition.ndim): |
| 963 | arr.append(index[index_ele, :]) |
| 964 | ret = tuple(arr) |
| 965 | else: |
| 966 | ret = transpose(index, (1, 0)) |
| 967 | return ret |
| 968 | |
| 969 | |
| 970 | def swapaxes(inp: Tensor, axis1: int, axis2: int) -> Tensor: |