Esempio n. 1
0
    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)
Esempio n. 2
0
    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)