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 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
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
