| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # pylint: disable=invalid-name, too-many-arguments, too-many-nested-blocks |
| """Argwhere operator""" |
| from tvm.te import hybrid |
| |
| |
| @hybrid.script |
| def hybrid_argwhere_1d(output_shape, condition): |
| """Find the indices of elements of a 1-D tensor that are non-zero. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| 1-D tensor with boolean values. |
| |
| Returns |
| ------- |
| out : tvm.te.Tensor |
| Indices of non-zero elements. |
| """ |
| a = output_tensor(output_shape, "int32") |
| a1 = condition.shape[0] |
| valid_index = 0 |
| for i1 in range(a1): |
| if condition[i1] != 0: |
| a[valid_index, 0] = i1 |
| valid_index += 1 |
| return a |
| |
| |
| @hybrid.script |
| def hybrid_argwhere_2d(output_shape, condition): |
| """Find the indices of elements of a 2-D tensor that are non-zero. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| 2-D tensor with boolean values. |
| |
| Returns |
| ------- |
| out : tvm.te.Tensor |
| Indices of non-zero elements. |
| """ |
| a = output_tensor(output_shape, "int32") |
| a1 = condition.shape[0] |
| a2 = condition.shape[1] |
| valid_index = 0 |
| for i1 in range(a1): |
| for i2 in range(a2): |
| if condition[i1, i2] != 0: |
| a[valid_index, 0] = i1 |
| a[valid_index, 1] = i2 |
| valid_index += 1 |
| return a |
| |
| |
| @hybrid.script |
| def hybrid_argwhere_3d(output_shape, condition): |
| """Find the indices of elements of a 3-D tensor that are non-zero. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| 3-D tensor with boolean values. |
| |
| Returns |
| ------- |
| out : tvm.te.Tensor |
| Indices of non-zero elements. |
| """ |
| a = output_tensor(output_shape, "int32") |
| a1 = condition.shape[0] |
| a2 = condition.shape[1] |
| a3 = condition.shape[2] |
| valid_index = 0 |
| for i1 in range(a1): |
| for i2 in range(a2): |
| for i3 in range(a3): |
| if condition[i1, i2, i3] != 0: |
| a[valid_index, 0] = i1 |
| a[valid_index, 1] = i2 |
| a[valid_index, 2] = i3 |
| valid_index += 1 |
| return a |
| |
| |
| @hybrid.script |
| def hybrid_argwhere_4d(output_shape, condition): |
| """Find the indices of elements of a 4-D tensor that are non-zero. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| 4-D tensor with boolean values. |
| |
| Returns |
| ------- |
| out : tvm.te.Tensor |
| Indices of non-zero elements. |
| """ |
| a = output_tensor(output_shape, "int32") |
| a1 = condition.shape[0] |
| a2 = condition.shape[1] |
| a3 = condition.shape[2] |
| a4 = condition.shape[3] |
| valid_index = 0 |
| for i1 in range(a1): |
| for i2 in range(a2): |
| for i3 in range(a3): |
| for i4 in range(a4): |
| if condition[i1, i2, i3, i4] != 0: |
| a[valid_index, 0] = i1 |
| a[valid_index, 1] = i2 |
| a[valid_index, 2] = i3 |
| a[valid_index, 3] = i4 |
| valid_index += 1 |
| return a |
| |
| |
| @hybrid.script |
| def hybrid_argwhere_5d(output_shape, condition): |
| """Find the indices of elements of a 5-D tensor that are non-zero. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| 5-D tensor with boolean values. |
| |
| Returns |
| ------- |
| out : tvm.te.Tensor |
| Indices of non-zero elements. |
| """ |
| a = output_tensor(output_shape, "int32") |
| a1 = condition.shape[0] |
| a2 = condition.shape[1] |
| a3 = condition.shape[2] |
| a4 = condition.shape[3] |
| a5 = condition.shape[4] |
| valid_index = 0 |
| for i1 in range(a1): |
| for i2 in range(a2): |
| for i3 in range(a3): |
| for i4 in range(a4): |
| for i5 in range(a5): |
| if condition[i1, i2, i3, i4, i5] != 0: |
| a[valid_index, 0] = i1 |
| a[valid_index, 1] = i2 |
| a[valid_index, 2] = i3 |
| a[valid_index, 3] = i4 |
| a[valid_index, 4] = i5 |
| valid_index += 1 |
| return a |
| |
| |
| def argwhere(output_shape, condition): |
| """Find the indices of elements of a tensor that are non-zero. |
| |
| Parameters |
| ---------- |
| condition : tvm.te.Tensor |
| Tensor with boolean values. |
| |
| Returns |
| ------- |
| out : tvm.te.Tensor |
| Indices of non-zero elements. |
| """ |
| if len(condition.shape) == 1: |
| return hybrid_argwhere_1d(output_shape.shape, condition) |
| if len(condition.shape) == 2: |
| return hybrid_argwhere_2d(output_shape.shape, condition) |
| if len(condition.shape) == 3: |
| return hybrid_argwhere_3d(output_shape.shape, condition) |
| if len(condition.shape) == 4: |
| return hybrid_argwhere_4d(output_shape.shape, condition) |
| if len(condition.shape) == 5: |
| return hybrid_argwhere_5d(output_shape.shape, condition) |
| raise ValueError("Does not support rank higher than 5 in argwhere") |
