Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: support arbitrary ufuncs with respect to NEP-50 #2799

Merged
merged 7 commits into from
Nov 6, 2023
Merged

fix: support arbitrary ufuncs with respect to NEP-50 #2799

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
75b9506
fix: support arbtirary dtypes
agoose77 Nov 6, 2023
548c890
test: add simple test
agoose77 Nov 6, 2023
a48033e
test: prepare for Windows
agoose77 Nov 6, 2023
0b9b9e0
test: fix test after investigating windows
agoose77 Nov 6, 2023
e3f4a4b
test: anticipate the system default int type
agoose77 Nov 6, 2023
60ccf30
test: does this work?
agoose77 Nov 6, 2023
3092be6
test: only 32-bit on 32-bit platforms?
agoose77 Nov 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 63 additions & 60 deletions src/awkward/_nplikes/array_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,6 @@
import math
from functools import lru_cache

import numpy
import packaging.version

from awkward._nplikes.numpy_like import (
ArrayLike,
IndexType,
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
177 changes: 90 additions & 87 deletions src/awkward/_nplikes/typetracer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
from typing import Callable

import numpy
import packaging.version

import awkward as ak
from awkward._nplikes.dispatch import register_nplike
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down
Loading