awkward
awkward copied to clipboard
scikit-hep
fix: return ufunc as it is for user defined `vectorize`d funcs
Codecov Report
All modified and coverable lines are covered by tests :white_check_mark:
Project coverage is 81.94%. Comparing base (
b749e49) to head (0c79601). Report is 30 commits behind head on main.
Additional details and impacted files
| Files | Coverage Δ | |
|---|---|---|
| src/awkward/_connect/jax/__init__.py | 100.00% <100.00%> (ø) |
![codecov[bot] avatar](https://avatars.githubusercontent.com/in/254?v=4 "Published by a pub.dev verified publisher"){kind=small}
@Saransh-cpp this looks good to me! I anticipate problems if the custom ufunc doesn't understand JAX arrays, but I suppose for now we can try the ufunc to see what happens. In future, we might want some kind of mechanism to associate e.g. a numba ufunc with a JAX-aware overload, but that can be a later step.
I anticipate problems if the custom ufunc doesn't understand JAX arrays
I see, this makes sense. Looking at this again, shouldn't this be on the user side? If I understand correctly, shouldn't a user ensure that their functions are compatible with JAX arrays if they are planning to use the jax backend?
I anticipate problems if the custom ufunc doesn't understand JAX arrays
I see, this makes sense. Looking at this again, shouldn't this be on the user side? If I understand correctly, shouldn't a user ensure that their functions are compatible with JAX arrays if they are planning to use the jax backend?
Yes I think so. As it stands, I'm not sure what that mechanism would be. It would be nice if a user can leverage someone-else's code, and patch their ufuncs to make JAX work, for example.
Yes I think so. As it stands, I'm not sure what that mechanism would be. It would be nice if a user can leverage someone-else's code, and patch their ufuncs to make JAX work, for example.
Oh I see, yes that would be nice too. I can look into this if any issues pop up in the future. Thanks!
The fix works! Jax and numba might still clash but that would be on the user side and not on the awkward side. But as mentioned by @agoose77, I'll look into providing users with an awkward-fix (maybe a decorator?) for their functions if we get any complaints.