1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17# pylint: disable=redefined-builtin,consider-using-enumerate,no-member
18"""Reduce operators"""
19from __future__ import absolute_import as _abs
20from . import cpp
21
22def _get_real_axis(ndim, axis):
23    if axis is None:
24        real_axis = list(range(ndim))
25    else:
26        if isinstance(axis, int):
27            axis = [axis]
28        else:
29            assert isinstance(axis, (list, tuple))
30        real_axis = []
31        for ele in axis:
32            if ele < 0:
33                ele += ndim
34            if ele >= ndim:
35                raise ValueError(
36                    "{} exceeds the maximum dimension {}. Received axis={}".format(ele, ndim, axis))
37            real_axis.append(ele)
38        real_axis.sort()
39        real_axis = list(set(real_axis))  # Remove the duplicates
40    return real_axis
41
42
43def sum(data, axis=None, keepdims=False):
44    """Sum of array elements over a given axis or a list of axes
45
46    Parameters
47    ----------
48    data : tvm.Tensor
49        The input tvm tensor
50
51    axis : None or int or tuple of int
52        Axis or axes along which a sum is performed.
53        The default, axis=None, will sum all of the elements of the input array.
54        If axis is negative it counts from the last to the first axis.
55
56    keepdims : bool
57        If this is set to True, the axes which are reduced are left in the result as dimensions
58        with size one.
59        With this option, the result will broadcast correctly against the input array.
60
61    Returns
62    -------
63    ret : tvm.Tensor
64    """
65    return cpp.sum(data, axis, keepdims)
66
67
68def all(data, axis=None, keepdims=False):
69    """Logical AND of array elements over a given axis or a list of axes
70
71    Parameters
72    ----------
73    data : tvm.Tensor
74        The input tvm boolean tensor
75
76    axis : None or int or tuple of int
77        Axis or axes along which a logical AND is performed.
78        The default, axis=None, will perform logical AND over all elements of the input array.
79        If axis is negative it counts from the last to the first axis.
80
81    keepdims : bool
82        If this is set to True, the axes which are reduced are left in the result as dimensions
83        with size one.
84        With this option, the result will broadcast correctly against the input array.
85
86    Returns
87    -------
88    ret : tvm.Tensor
89    """
90    return cpp.all(data, axis, keepdims)
91
92
93def any(data, axis=None, keepdims=False):
94    """Logical OR of array elements over a given axis or a list of axes
95
96    Parameters
97    ----------
98    data : tvm.Tensor
99        The input tvm boolean tensor
100
101    axis : None or int or tuple of int
102        Axis or axes along which a logical OR is performed.
103        The default, axis=None, will perform logical OR over all elements of the input array.
104        If axis is negative it counts from the last to the first axis.
105
106    keepdims : bool
107        If this is set to True, the axes which are reduced are left in the result as dimensions
108        with size one.
109        With this option, the result will broadcast correctly against the input array.
110
111    Returns
112    -------
113    ret : tvm.Tensor
114    """
115    return cpp.any(data, axis, keepdims)
116
117
118def max(data, axis=None, keepdims=False):
119    """Maximum of array elements over a given axis or a list of axes
120
121    Parameters
122    ----------
123    data : tvm.Tensor
124        The input tvm tensor
125
126    axis : None or int or tuple of int
127        Axis or axes along which the max operation is performed.
128        The default, axis=None, will find the max element from all of the elements of the input
129        array. If axis is negative it counts from the last to the first axis.
130
131    keepdims : bool
132        If this is set to True, the axes which are reduced are left in the result as dimensions
133        with size one.
134        With this option, the result will broadcast correctly against the input array.
135
136    Returns
137    -------
138    ret : tvm.Tensor
139    """
140    return cpp.max(data, axis, keepdims)
141
142
143def min(data, axis=None, keepdims=False):
144    """Minimum of array elements over a given axis or a list of axes
145
146    Parameters
147    ----------
148    data : tvm.Tensor
149        The input tvm tensor
150
151    axis : None or int or tuple of int
152        Axis or axes along which a minimum operation is performed.
153        The default, axis=None, will find the minimum element from all of the elements of the
154        input array. If axis is negative it counts from the last to the first axis.
155
156    keepdims : bool
157        If this is set to True, the axes which are reduced are left in the result as dimensions
158        with size one.
159        With this option, the result will broadcast correctly against the input array.
160
161    Returns
162    -------
163    ret : tvm.Tensor
164    """
165    return cpp.min(data, axis, keepdims)
166
167
168def argmax(data, axis=None, keepdims=False):
169    """Returns the indices of the maximum values along an axis.
170
171    Parameters
172    ----------
173    data : tvm.Tensor
174        The input tvm tensor
175
176    axis : None or int or tuple of int
177        Axis or axes along which a argmax operation is performed.
178        The default, axis=None, will find the indices of the maximum element of the elements of
179        the input array. If axis is negative it counts from the last to the first axis.
180
181    keepdims : bool
182        If this is set to True, the axes which are reduced are left in the result as dimensions
183        with size one.
184        With this option, the result will broadcast correctly against the input array.
185
186    Returns
187    -------
188    ret : tvm.Tensor
189    """
190    return cpp.argmax(data, axis, keepdims)
191
192
193def argmin(data, axis=None, keepdims=False):
194    """Returns the indices of the minimum values along an axis.
195
196    Parameters
197    ----------
198    data : tvm.Tensor
199        The input tvm tensor
200
201    axis : None or int or tuple of int
202        Axis or axes along which a argmin operation is performed.
203        The default, axis=None, will find the indices of minimum element all of the elements of
204        the input array. If axis is negative it counts from the last to the first axis.
205
206    keepdims : bool
207        If this is set to True, the axes which are reduced are left in the result as dimensions
208        with size one.
209        With this option, the result will broadcast correctly against the input array.
210
211    Returns
212    -------
213    ret : tvm.Tensor
214    """
215    return cpp.argmin(data, axis, keepdims)
216
217
218def prod(data, axis=None, keepdims=False):
219    """Product of array elements over a given axis or a list of axes
220
221    Parameters
222    ----------
223    data : tvm.Tensor
224        The input tvm tensor
225
226    axis : None or int or tuple of int
227        Axis or axes along which a prod operation is performed.
228        The default, axis=None, will get the prod element over all of the elements of the
229        input array. If axis is negative it counts from the last to the first axis.
230
231    keepdims : bool
232        If this is set to True, the axes which are reduced are left in the result as dimensions
233        with size one.
234        With this option, the result will broadcast correctly against the input array.
235
236    Returns
237    -------
238    ret : tvm.Tensor
239    """
240    return cpp.prod(data, axis, keepdims)
241