diff --git a/src/awkward/_nplikes/array_module.py b/src/awkward/_nplikes/array_module.py index 6e02480c5f..1be31cea19 100644 --- a/src/awkward/_nplikes/array_module.py +++ b/src/awkward/_nplikes/array_module.py @@ -4,9 +4,6 @@ import math from functools import lru_cache -import numpy -import packaging.version - from awkward._nplikes.numpy_like import ( ArrayLike, IndexType, @@ -23,9 +20,6 @@ from numpy.typing import DTypeLike np = NumpyMetadata.instance() -NUMPY_HAS_NEP_50 = packaging.version.Version( - numpy.__version__ -) >= packaging.version.Version("1.24") @lru_cache @@ -198,62 +192,71 @@ def searchsorted( return self._module.searchsorted(x, values, side=side, sorter=sorter) ############################ manipulation + def apply_ufunc( + self, + ufunc: UfuncLike, + method: str, + args: list[Any], + kwargs: dict[str, Any] | None = None, + ) -> ArrayLikeT | tuple[ArrayLikeT, ...]: + if method != "__call__" or len(args) == 0: + raise NotImplementedError + + if hasattr(ufunc, "resolve_dtypes"): + return self._apply_ufunc_nep_50(ufunc, method, args, kwargs) + else: + return self._apply_ufunc_legacy(ufunc, method, args, kwargs) # Does NumPy support value-less ufunc resolution? - if NUMPY_HAS_NEP_50: - - def apply_ufunc( - self, - ufunc: UfuncLike, - method: str, - args: list[Any], - kwargs: dict[str, Any] | None = None, - ) -> ArrayLikeT | tuple[ArrayLikeT]: - # Determine input argument dtypes - input_arg_dtypes = [getattr(obj, "dtype", type(obj)) for obj in args] - # Resolve these for the given ufunc - arg_dtypes = tuple(input_arg_dtypes + [None] * ufunc.nout) - resolved_dtypes = ufunc.resolve_dtypes(arg_dtypes) - # Interpret the arguments under these dtypes, converting scalars to length-1 arrays - resolved_args = [ - cast("ArrayLikeT", self.asarray(arg, dtype=dtype)) - for arg, dtype in zip(args, resolved_dtypes) - ] - # Broadcast to ensure all-scalar or all-nd-array - broadcasted_args = self.broadcast_arrays(*resolved_args) - # Allow other nplikes to replace implementation - impl = self.prepare_ufunc(ufunc) - # Compute the result - return impl(*broadcasted_args, **(kwargs or {})) - - else: - # Otherwise, perform default NumPy coercion (value-dependent) - def apply_ufunc( - self, - ufunc: UfuncLike, - method: str, - args: list[Any], - kwargs: dict[str, Any] | None = None, - ) -> ArrayLikeT | tuple[ArrayLikeT]: - # Convert np.generic to scalar arrays - resolved_args = [ - cast( - "ArrayLikeT", - self.asarray( - arg, dtype=arg.dtype if hasattr(arg, "dtype") else None - ), - ) - for arg in args - ] - broadcasted_args = self.broadcast_arrays(*resolved_args) - # Choose the broadcasted argument if it wasn't a Python scalar - non_generic_value_promoted_args = [ - y if hasattr(x, "ndim") else x for x, y in zip(args, broadcasted_args) - ] - # Allow other nplikes to replace implementation - impl = self.prepare_ufunc(ufunc) - # Compute the result - return impl(*non_generic_value_promoted_args, **(kwargs or {})) + def _apply_ufunc_nep_50( + self, + ufunc: UfuncLike, + method: str, + args: list[Any], + kwargs: dict[str, Any] | None = None, + ) -> ArrayLikeT | tuple[ArrayLikeT]: + # Determine input argument dtypes + input_arg_dtypes = [getattr(obj, "dtype", type(obj)) for obj in args] + # Resolve these for the given ufunc + arg_dtypes = tuple(input_arg_dtypes + [None] * ufunc.nout) + resolved_dtypes = ufunc.resolve_dtypes(arg_dtypes) + # Interpret the arguments under these dtypes, converting scalars to length-1 arrays + resolved_args = [ + cast("ArrayLikeT", self.asarray(arg, dtype=dtype)) + for arg, dtype in zip(args, resolved_dtypes) + ] + # Broadcast to ensure all-scalar or all-nd-array + broadcasted_args = self.broadcast_arrays(*resolved_args) + # Allow other nplikes to replace implementation + impl = self.prepare_ufunc(ufunc) + # Compute the result + return impl(*broadcasted_args, **(kwargs or {})) + + # Otherwise, perform default NumPy coercion (value-dependent) + def _apply_ufunc_legacy( + self, + ufunc: UfuncLike, + method: str, + args: list[Any], + kwargs: dict[str, Any] | None = None, + ) -> ArrayLikeT | tuple[ArrayLikeT]: + # Convert np.generic to scalar arrays + resolved_args = [ + cast( + "ArrayLikeT", + self.asarray(arg, dtype=arg.dtype if hasattr(arg, "dtype") else None), + ) + for arg in args + ] + broadcasted_args = self.broadcast_arrays(*resolved_args) + # Choose the broadcasted argument if it wasn't a Python scalar + non_generic_value_promoted_args = [ + y if hasattr(x, "ndim") else x for x, y in zip(args, broadcasted_args) + ] + # Allow other nplikes to replace implementation + impl = self.prepare_ufunc(ufunc) + # Compute the result + return impl(*non_generic_value_promoted_args, **(kwargs or {})) def broadcast_arrays(self, *arrays: ArrayLikeT) -> list[ArrayLikeT]: assert not any(isinstance(x, PlaceholderArray) for x in arrays) diff --git a/src/awkward/_nplikes/typetracer.py b/src/awkward/_nplikes/typetracer.py index 0f4155413f..2e98018a56 100644 --- a/src/awkward/_nplikes/typetracer.py +++ b/src/awkward/_nplikes/typetracer.py @@ -6,7 +6,6 @@ from typing import Callable import numpy -import packaging.version import awkward as ak from awkward._nplikes.dispatch import register_nplike @@ -45,9 +44,6 @@ np = NumpyMetadata.instance() -NUMPY_HAS_NEP_50 = packaging.version.Version( - numpy.__version__ -) >= packaging.version.Version("1.24") def is_unknown_length(array: Any) -> bool: @@ -484,9 +480,6 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): # ) kwargs.pop("out", None) - if method != "__call__" or len(inputs) == 0: - raise NotImplementedError - if len(kwargs) > 0: raise ValueError("TypeTracerArray does not support kwargs for ufuncs") return self.nplike.apply_ufunc(ufunc, method, inputs, kwargs) @@ -530,89 +523,99 @@ class TypeTracer(NumpyLike[TypeTracerArray]): is_eager: Final = True supports_structured_dtypes: Final = True - if NUMPY_HAS_NEP_50: - - def apply_ufunc( - self, - ufunc: UfuncLike, - method: str, - args: Sequence[Any], - kwargs: dict[str, Any] | None = None, - ) -> TypeTracerArray | tuple[TypeTracerArray, ...]: - for x in args: - try_touch_data(x) - - # Unwrap options, assume they don't occur - args = [x.content if isinstance(x, MaybeNone) else x for x in args] - # Determine input argument dtypes - input_arg_dtypes = [getattr(obj, "dtype", type(obj)) for obj in args] - # Resolve these for the given ufunc - arg_dtypes = tuple(input_arg_dtypes + [None] * ufunc.nout) - resolved_dtypes = ufunc.resolve_dtypes(arg_dtypes) - # Interpret the arguments under these dtypes - resolved_args = [ - self.asarray(arg, dtype=dtype) - for arg, dtype in zip(args, resolved_dtypes) - ] - # Broadcast to ensure all-scalar or all-nd-array - broadcasted_args = self.broadcast_arrays(*resolved_args) - broadcasted_shape = broadcasted_args[0].shape - result_dtypes = resolved_dtypes[ufunc.nin :] - - if len(result_dtypes) == 1: - return TypeTracerArray._new(result_dtypes[0], shape=broadcasted_shape) - else: - return tuple( - TypeTracerArray._new(dtype, shape=broadcasted_shape) - for dtype in result_dtypes - ) + def apply_ufunc( + self, + ufunc: UfuncLike, + method: str, + args: list[Any], + kwargs: dict[str, Any] | None = None, + ) -> TypeTracerArray | tuple[TypeTracerArray, ...]: + if method != "__call__" or len(args) == 0: + raise NotImplementedError - else: + if hasattr(ufunc, "resolve_dtypes"): + return self._apply_ufunc_nep_50(ufunc, method, args, kwargs) + else: + return self._apply_ufunc_legacy(ufunc, method, args, kwargs) - def apply_ufunc( - self, - ufunc: UfuncLike, - method: str, - args: Sequence[Any], - kwargs: dict[str, Any] | None = None, - ) -> TypeTracerArray | tuple[TypeTracerArray, ...]: - for x in args: - try_touch_data(x) - - # Unwrap options, assume they don't occur - args = [x.content if isinstance(x, MaybeNone) else x for x in args] - # Convert np.generic to scalar arrays - resolved_args = [ - self.asarray(arg, dtype=arg.dtype if hasattr(arg, "dtype") else None) - for arg in args - ] - # Broadcast all inputs together - broadcasted_args = self.broadcast_arrays(*resolved_args) - broadcasted_shape = broadcasted_args[0].shape - # Choose the broadcasted argument if it wasn't a Python scalar - non_generic_value_promoted_args = [ - y if hasattr(x, "ndim") else x for x, y in zip(args, broadcasted_args) - ] - # Build proxy (empty) arrays - proxy_args = [ - (numpy.empty(0, dtype=x.dtype) if hasattr(x, "dtype") else x) - for x in non_generic_value_promoted_args - ] - # Determine result dtype from proxy call - proxy_result = ufunc(*proxy_args, **(kwargs or {})) - if ufunc.nout == 1: - result_dtypes = [proxy_result.dtype] - else: - assert isinstance(proxy_result, tuple) - result_dtypes = [x.dtype for x in proxy_result] + def _apply_ufunc_nep_50( + self, + ufunc: UfuncLike, + method: str, + args: Sequence[Any], + kwargs: dict[str, Any] | None = None, + ) -> TypeTracerArray | tuple[TypeTracerArray, ...]: + for x in args: + try_touch_data(x) - if len(result_dtypes) == 1: - return TypeTracerArray._new(result_dtypes[0], shape=broadcasted_shape) - else: - return tuple( - TypeTracerArray._new(dtype, shape=broadcasted_shape) - for dtype in result_dtypes - ) + # Unwrap options, assume they don't occur + args = [x.content if isinstance(x, MaybeNone) else x for x in args] + # Determine input argument dtypes + input_arg_dtypes = [getattr(obj, "dtype", type(obj)) for obj in args] + # Resolve these for the given ufunc + arg_dtypes = tuple(input_arg_dtypes + [None] * ufunc.nout) + resolved_dtypes = ufunc.resolve_dtypes(arg_dtypes) + # Interpret the arguments under these dtypes + resolved_args = [ + self.asarray(arg, dtype=dtype) for arg, dtype in zip(args, resolved_dtypes) + ] + # Broadcast to ensure all-scalar or all-nd-array + broadcasted_args = self.broadcast_arrays(*resolved_args) + broadcasted_shape = broadcasted_args[0].shape + result_dtypes = resolved_dtypes[ufunc.nin :] + + if len(result_dtypes) == 1: + return TypeTracerArray._new(result_dtypes[0], shape=broadcasted_shape) + else: + return tuple( + TypeTracerArray._new(dtype, shape=broadcasted_shape) + for dtype in result_dtypes + ) + + def _apply_ufunc_legacy( + self, + ufunc: UfuncLike, + method: str, + args: Sequence[Any], + kwargs: dict[str, Any] | None = None, + ) -> TypeTracerArray | tuple[TypeTracerArray, ...]: + for x in args: + try_touch_data(x) + + # Unwrap options, assume they don't occur + args = [x.content if isinstance(x, MaybeNone) else x for x in args] + # Convert np.generic to scalar arrays + resolved_args = [ + self.asarray(arg, dtype=arg.dtype if hasattr(arg, "dtype") else None) + for arg in args + ] + # Broadcast all inputs together + broadcasted_args = self.broadcast_arrays(*resolved_args) + broadcasted_shape = broadcasted_args[0].shape + # Choose the broadcasted argument if it wasn't a Python scalar + non_generic_value_promoted_args = [ + y if hasattr(x, "ndim") else x for x, y in zip(args, broadcasted_args) + ] + # Build proxy (empty) arrays + proxy_args = [ + (numpy.empty(0, dtype=x.dtype) if hasattr(x, "dtype") else x) + for x in non_generic_value_promoted_args + ] + # Determine result dtype from proxy call + proxy_result = ufunc(*proxy_args, **(kwargs or {})) + if ufunc.nout == 1: + result_dtypes = [proxy_result.dtype] + else: + assert isinstance(proxy_result, tuple) + result_dtypes = [x.dtype for x in proxy_result] + + if len(result_dtypes) == 1: + return TypeTracerArray._new(result_dtypes[0], shape=broadcasted_shape) + else: + return tuple( + TypeTracerArray._new(dtype, shape=broadcasted_shape) + for dtype in result_dtypes + ) def _axis_is_valid(self, axis: int, ndim: int) -> bool: if axis < 0: diff --git a/tests/test_2799_numba_ufunc_resolution.py b/tests/test_2799_numba_ufunc_resolution.py new file mode 100644 index 0000000000..63bc586c73 --- /dev/null +++ b/tests/test_2799_numba_ufunc_resolution.py @@ -0,0 +1,46 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +import numpy as np +import packaging.version +import pytest + +import awkward as ak + +numba = pytest.importorskip("numba") + +NUMBA_HAS_NEP_50 = packaging.version.parse( + numba.__version__ +) >= packaging.version.Version("0.59.0") + + +@pytest.mark.skipif(not NUMBA_HAS_NEP_50, reason="Numba does not have NEP-50 support") +def test_numba_ufunc_nep_50(): + raise NotImplementedError + + @numba.vectorize + def add(x, y): + return x + y + + array = ak.values_astype([[1, 2, 3], [4]], np.int8) + + # FIXME: what error will Numba throw here for an out-of-bounds integer? + with pytest.warns(FutureWarning, match=r"not create a writeable array"): + result = add(array, np.int16(np.iinfo(np.int8).max + 1)) + + flattened = ak.to_numpy(ak.flatten(result)) + assert flattened.dtype == np.dtype(np.int8) + + +@pytest.mark.skipif(NUMBA_HAS_NEP_50, reason="Numba has NEP-50 support") +def test_numba_ufunc_legacy(): + @numba.vectorize + def add(x, y): + return x + y + + array = ak.values_astype([[1, 2, 3], [4]], np.int8) + with pytest.warns(FutureWarning, match=r"not create a writeable array"): + result = add(array, np.int16(np.iinfo(np.int8).max + 1)) + + flattened = ak.to_numpy(ak.flatten(result)) + # Seems like Numba chooses int64 here unless a 32-bit platform + assert flattened.dtype == np.dtype(np.int32 if ak._util.bits32 else np.int64)