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.

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

Add __array_function__ to dispatch NumPy funcs to dask-awkward #520

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

supercoder-dev
Copy link

@supercoder-dev supercoder-dev commented Jun 26, 2024

Add NumPy function dispatching to dask-awkward

This PR adds an array_function method to the Array class, allowing NumPy functions to be automatically redirected to their dask-awkward counterparts. It covers 38 common functions, making it easier to use NumPy-style operations with dask-awkward arrays.

Copy link
Collaborator

@martindurant martindurant left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could use tests.

However: isn't the purpose of array_function to return a numpy view on the data rather than (as with ufuncs) to apply the numpy idea to producing an array of our own type? Maybe I misunderstand.

if func.__name__ in supported_funcs:
return supported_funcs[func.__name__](*args, **kwargs)
else:
return NotImplemented
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
return NotImplemented
return NotImplemented(str(func))

import dask_awkward as dak

# List of supported functions
supported_funcs = {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mildly prefer this as a global constant rather than in the function

@lgray
Copy link
Collaborator

lgray commented Jun 26, 2024

I'd be a bit careful with this kind of dispatch. The numpy functions can have subtly different behaviors from their awkward counterparts. This mapping isn't very well defined for the arbitrary numpy functions (but is for ufuncs).

It would almost be better to suggest converting to dask.array and error out since there the expected behavior is much more clear.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants