def tensordot(context, 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 = context.promote_types(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 PyArray_Empty(args, name='PyArray_Empty'):
nd, shape, dtype, fortran = args
return_type = typesystem.array(dtype, nd)
signature = return_type(
int_, # nd
npy_intp.pointer(), # shape
object_, # dtype
int_) # fortran
return MultiArrayAPINode(name, signature, args)
def array_from_type(type):
if type.is_array:
return type
elif type.is_tuple or type.is_list:
dtype = array_from_type(type.base_type)
if dtype.is_array:
return dtype.add('ndim', type.ndim + 1)
elif not type.is_object:
return typesystem.array(dtype=type, ndim=0)
return object_
def reduce_(a, axis, dtype, out, static_dtype=None):
if out is not None:
return out
dtype_type = _dtype(a, dtype, static_dtype)
if axis is None:
# Return the scalar type
return dtype_type
if dtype_type:
# Handle the axis parameter
if axis.is_tuple and axis.is_sized:
# axis=(tuple with a constant size)
return typesystem.array(dtype_type, a.ndim - axis.size)
elif axis.is_int:
# axis=1
return typesystem.array(dtype_type, a.ndim - 1)
else:
# axis=(something unknown)
return object_
def dot(context, 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 = context.promote_types(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 empty(shape, dtype, order):
if shape is None:
return None
dtype = get_dtype(dtype, float64)
if shape.is_int:
ndim = 1
elif shape.is_tuple or shape.is_list:
ndim = shape.size
else:
return None
return typesystem.array(dtype.dtype, ndim)
def empty_like(a, dtype, order):
"Parse the result type for np.empty_like calls"
if a is None:
return None
if a.is_array:
if dtype:
dtype_type = get_dtype(dtype)
if dtype_type is None:
return a
dtype = dtype_type.dtype
else:
dtype = a.dtype
return typesystem.array(dtype, a.ndim)
def inner(context, a, b):
lhs_type = promote_to_array(a)
rhs_type = promote_to_array(b)
dtype = context.promote_types(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 empty_like(a, dtype, order):
"Parse the result type for np.empty_like calls"
if a is None:
return None
if a.is_array:
if dtype:
dtype_type = get_dtype(dtype)
if dtype_type is None:
return a
dtype = dtype_type.dtype
else:
dtype = a.dtype
return typesystem.array(dtype, a.ndim)
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 empty(shape, dtype, order):
if shape is None:
return None
dtype = get_dtype(dtype, float64)
if dtype is None:
return object_
if shape.is_int:
ndim = 1
elif shape.is_tuple or shape.is_list:
ndim = shape.size
else:
return None
return typesystem.array(dtype.dtype, ndim)
def infer(context, *args, **kwargs):
if len(args) != ufunc.nin:
return object_
# Find the right ufunc signature
argtypes = [type.dtype if type.is_array else type for type in args]
result_type = ufunc_infer.infer(context, argtypes)
if result_type is None:
return object_
# Determine output ndim
ndim = 0
for argtype in args:
if argtype.is_array:
ndim = max(argtype.ndim, ndim)
return typesystem.array(result_type, ndim)
def infer(context, *args, **kwargs):
if len(args) != ufunc.nin:
return object_
# Find the right ufunc signature
argtypes = [type.dtype if type.is_array else type for type in args]
result_type = ufunc_infer.infer(context, argtypes)
if result_type is None:
return object_
# Determine output ndim
ndim = 0
for argtype in args:
if argtype.is_array:
ndim = max(argtype.ndim, ndim)
return typesystem.array(result_type, ndim)
