def __call__(self, *args, **kwargs):
# If compilation is disabled OR it is NOT a dynamic gufunc
# call the underlying gufunc
if self._frozen or not self.is_dynamic:
return self.ufunc(*args, **kwargs)
elif "out" in kwargs:
# If "out" argument is supplied
args += (kwargs.pop("out"), )
if self._num_args_match(*args) is False:
# It is not allowed to call a dynamic gufunc without
# providing all the arguments
# see: https://github.com/numba/numba/pull/5938#discussion_r506429392 # noqa: E501
msg = (
f"Too few arguments for function '{self.__name__}'. "
"Note that the pattern `out = gufunc(Arg1, Arg2, ..., ArgN)` "
"is not allowed. Use `gufunc(Arg1, Arg2, ..., ArgN, out) "
"instead.")
raise TypeError(msg)
# at this point we know the gufunc is a dynamic one
ewise = self._get_ewise_dtypes(args)
if not (self.ufunc and ufunc_find_matching_loop(self.ufunc, ewise)):
sig = self._get_signature(*args)
self.add(sig)
self.build_ufunc()
return self.ufunc(*args, **kwargs)
def __init__(self, context, builder, outer_sig):
super(_KernelImpl, self).__init__(context, builder, outer_sig)
loop = ufunc_find_matching_loop(
ufunc, outer_sig.args + (outer_sig.return_type, ))
self.fn = ufunc_db.get_ufunc_info(ufunc).get(loop.ufunc_sig)
self.inner_sig = typing.signature(*(loop.outputs + loop.inputs))
if self.fn is None:
msg = "Don't know how to lower ufunc '{0}' for loop '{1}'"
raise NotImplementedError(msg.format(ufunc.__name__, loop))
def __init__(self, context, builder, outer_sig):
super(_KernelImpl, self).__init__(context, builder, outer_sig)
loop = ufunc_find_matching_loop(
ufunc, outer_sig.args + tuple(_unpack_output_types(ufunc, outer_sig)))
self.fn = context.get_ufunc_info(ufunc).get(loop.ufunc_sig)
self.inner_sig = _ufunc_loop_sig(loop.outputs, loop.inputs)
if self.fn is None:
msg = "Don't know how to lower ufunc '{0}' for loop '{1}'"
raise NotImplementedError(msg.format(ufunc.__name__, loop))
def generic(self, args, kws):
ufunc = self.ufunc
base_types, explicit_outputs, ndims, layout = self._handle_inputs(
ufunc, args, kws)
ufunc_loop = ufunc_find_matching_loop(ufunc, base_types)
if ufunc_loop is None:
raise TypingError("can't resolve ufunc {0} for types {1}".format(
ufunc.__name__, args))
# check if all the types involved in the ufunc loop are supported in this mode
if not supported_ufunc_loop(ufunc, ufunc_loop):
msg = "ufunc '{0}' using the loop '{1}' not supported in this mode"
raise TypingError(
msg=msg.format(ufunc.__name__, ufunc_loop.ufunc_sig))
# if there is any explicit output type, check that it is valid
explicit_outputs_np = [as_dtype(tp.dtype) for tp in explicit_outputs]
# Numpy will happily use unsafe conversions (although it will actually warn)
if not all(
np.can_cast(fromty, toty, 'unsafe')
for (fromty, toty
) in zip(ufunc_loop.numpy_outputs, explicit_outputs_np)):
msg = "ufunc '{0}' can't cast result to explicit result type"
raise TypingError(msg=msg.format(ufunc.__name__))
# A valid loop was found that is compatible. The result of type inference should
# be based on the explicit output types, and when not available with the type given
# by the selected NumPy loop
out = list(explicit_outputs)
implicit_output_count = ufunc.nout - len(explicit_outputs)
if implicit_output_count > 0:
# XXX this is sometimes wrong for datetime64 and timedelta64,
# as ufunc_find_matching_loop() doesn't do any type inference
ret_tys = ufunc_loop.outputs[-implicit_output_count:]
if ndims > 0:
assert layout is not None
ret_tys = [
types.Array(dtype=ret_ty, ndim=ndims, layout=layout)
for ret_ty in ret_tys
]
ret_tys = [
resolve_output_type(self.context, args, ret_ty)
for ret_ty in ret_tys
]
out.extend(ret_tys)
# note: although the previous code should support multiple return values, only one
# is supported as of now (signature may not support more than one).
# there is an check enforcing only one output
out.extend(args)
return signature(*out)
def find_ewise_function(self, ewise_types):
"""
Given a tuple of element-wise argument types, find a matching
signature in the dispatcher.
Return a 2-tuple containing the matching signature, and
compilation result. Will return two None's if no matching
signature was found.
"""
if self._frozen:
# If we cannot compile, coerce to the best matching loop
loop = numpy_support.ufunc_find_matching_loop(self, ewise_types)
if loop is None:
return None, None
ewise_types = tuple(loop.inputs + loop.outputs)[:len(ewise_types)]
for sig, cres in self._dispatcher.overloads.items():
if sig.args == ewise_types:
return sig, cres
return None, None
def generic(self, args, kws):
# First, strip optional types, ufunc loops are typed on concrete types
args = [x.type if isinstance(x, types.Optional) else x for x in args]
ufunc = self.ufunc
base_types, explicit_outputs, ndims, layout = self._handle_inputs(
ufunc, args, kws)
ufunc_loop = ufunc_find_matching_loop(ufunc, base_types)
if ufunc_loop is None:
raise TypingError("can't resolve ufunc {0} for types {1}".format(
ufunc.__name__, args))
# check if all the types involved in the ufunc loop are supported in this mode
if not supported_ufunc_loop(ufunc, ufunc_loop):
msg = "ufunc '{0}' using the loop '{1}' not supported in this mode"
raise TypingError(
msg=msg.format(ufunc.__name__, ufunc_loop.ufunc_sig))
# if there is any explicit output type, check that it is valid
explicit_outputs_np = [as_dtype(tp.dtype) for tp in explicit_outputs]
# Numpy will happily use unsafe conversions (although it will actually warn)
if not all(
np.can_cast(fromty, toty, 'unsafe')
for (fromty, toty
) in zip(ufunc_loop.numpy_outputs, explicit_outputs_np)):
msg = "ufunc '{0}' can't cast result to explicit result type"
raise TypingError(msg=msg.format(ufunc.__name__))
# A valid loop was found that is compatible. The result of type inference should
# be based on the explicit output types, and when not available with the type given
# by the selected NumPy loop
out = list(explicit_outputs)
implicit_output_count = ufunc.nout - len(explicit_outputs)
if implicit_output_count > 0:
# XXX this is sometimes wrong for datetime64 and timedelta64,
# as ufunc_find_matching_loop() doesn't do any type inference
ret_tys = ufunc_loop.outputs[-implicit_output_count:]
if ndims > 0:
assert layout is not None
# If either of the types involved in the ufunc operation have a
# __array_ufunc__ method then invoke the first such one to
# determine the output type of the ufunc.
array_ufunc_type = None
for a in args:
if hasattr(a, "__array_ufunc__"):
array_ufunc_type = a
break
output_type = types.Array
if array_ufunc_type is not None:
output_type = array_ufunc_type.__array_ufunc__(
ufunc, "__call__", *args, **kws)
if output_type is NotImplemented:
msg = (f"unsupported use of ufunc {ufunc} on "
f"{array_ufunc_type}")
# raise TypeError here because
# NumpyRulesArrayOperator.generic is capturing
# TypingError
raise NumbaTypeError(msg)
elif not issubclass(output_type, types.Array):
msg = (f"ufunc {ufunc} on {array_ufunc_type}"
f"cannot return non-array {output_type}")
# raise TypeError here because
# NumpyRulesArrayOperator.generic is capturing
# TypingError
raise TypeError(msg)
ret_tys = [
output_type(dtype=ret_ty, ndim=ndims, layout=layout)
for ret_ty in ret_tys
]
ret_tys = [
resolve_output_type(self.context, args, ret_ty)
for ret_ty in ret_tys
]
out.extend(ret_tys)
return _ufunc_loop_sig(out, args)
