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