def tensordot(typesystem, a, b, axes): '''Typing function for numpy.tensordot(). Defaults to Python object for any caller that isn't using the default argument to axes. Otherwise, it is similar to inner(), but subtracts four dimensions from the result instead of two. Without symbolic execution of the actual axes argument, this can't determine the number of axes to sum over, so it punts. This typing function could use an array type of unknown dimensionality, were one available. See: https://www.pivotaltracker.com/story/show/43687249 ''' lhs_type = array_from_object(a) rhs_type = array_from_object(b) if lhs_type.ndim < 1: raise error.NumbaError(a, 'First argument to numpy.tensordot() ' 'requires array of dimensionality >= 1.') elif rhs_type.ndim < 1: raise error.NumbaError(b, 'First argument to numpy.tensordot() ' 'requires array of dimensionality >= 1.') dtype = typesystem.promote(lhs_type.dtype, rhs_type.dtype) if axes is None: result_ndim = lhs_type.ndim + rhs_type.ndim - 4 if result_ndim < 0: raise error.NumbaError(a, 'Arguments to numpy.tensordot() should ' 'have combined dimensionality >= 4 (when ' 'axes argument is not specified).') result_type = typesystem.array(dtype, result_ndim) else: # XXX Issue warning to user? result_type = object_ return result_type
def minmax(typesystem, args, op): if len(args) < 2: return res = args[0] for arg in args[1:]: lhs_type = get_type(res) rhs_type = get_type(arg) res_type = typesystem.promote(lhs_type, rhs_type) if lhs_type != res_type: res = nodes.CoercionNode(res, res_type) if rhs_type != res_type: arg = nodes.CoercionNode(arg, res_type) lhs_temp = nodes.TempNode(res_type) rhs_temp = nodes.TempNode(res_type) res_temp = nodes.TempNode(res_type) lhs = lhs_temp.load(invariant=True) rhs = rhs_temp.load(invariant=True) expr = ast.IfExp(ast.Compare(lhs, [op], [rhs]), lhs, rhs) body = [ ast.Assign([lhs_temp.store()], res), ast.Assign([rhs_temp.store()], arg), ast.Assign([res_temp.store()], expr), ] res = nodes.ExpressionNode(body, res_temp.load(invariant=True)) return res
def dot(typesystem, a, b, out): "Resolve a call to np.dot()" if out is not None: return out lhs_type = promote_to_array(a) rhs_type = promote_to_array(b) dtype = typesystem.promote(lhs_type.dtype, rhs_type.dtype) dst_ndim = lhs_type.ndim + rhs_type.ndim - 2 result_type = typesystem.array(dtype, dst_ndim) return result_type
def inner(typesystem, a, b): lhs_type = promote_to_array(a) rhs_type = promote_to_array(b) dtype = typesystem.promote(lhs_type.dtype, rhs_type.dtype) if lhs_type.ndim == 0: result_ndim = rhs_type.ndim elif rhs_type.ndim == 0: result_ndim = lhs_type.ndim else: result_ndim = lhs_type.ndim + rhs_type.ndim - 2 if result_ndim == 0: result_type = dtype else: result_type = typesystem.array(dtype, result_ndim) return result_type
def vdot(typesystem, a, b): lhs_type = promote_to_array(a) rhs_type = promote_to_array(b) dtype = typesystem.promote(lhs_type.dtype, rhs_type.dtype) return dtype
def abs_(typesystem, node, x): argtype = typesystem.promote(long_, get_type(x)) dst_type = abstype(argtype) node.variable = Variable(dst_type) node.args = [nodes.CoercionNode(x, argtype)] return node