fix: return ufunc as it is for user defined vectorized funcs

[Saransh-cpp]

Fixes #2603

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

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%> (ø)

... and 3 files with indirect coverage changes

[agoose77]

@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.

[Saransh-cpp]

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?

[agoose77]

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.

[Saransh-cpp]

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!

[jpivarski]

What's the status of this one?

[Saransh-cpp]

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.

[jpivarski]

Okay, great! And it's such a simple change on the Awkward side, with a test, that I'll just merge it now. Thanks!