search

scikit-hep / awkward

Manipulate JSON-like data with NumPy-like idioms.
https://awkward-array.org
BSD 3-Clause "New" or "Revised" License
836 stars 87 forks source link

Auto-conversion to NumPy arrays within a Numbafied function #509

Open jpivarski opened 3 years ago

jpivarski commented 3 years ago

Inside a Numba-JITed function, users might want to call NumPy functions (np.*) on Awkward Arrays that happen to be 1D/simple/flat. Right now, they have to cast such an array with np.asarray and then call the NumPy function, which they don't have to do outside of Numba. And that's surprising. It may be possible to register Awkward Arrays (or just NumpyArrayType) as "array-like" in Numba, so that it does this conversion automatically.

jpivarski commented 3 years ago

This could go either way, "bug" or "feature," because it's highly expected, but I'm not certain how or how generally it can be implemented.

jpivarski commented 3 years ago

Questions about how to implement it on Numba's Discourse.

jpivarski commented 3 years ago

A few notes:

First, there's a @lower_cast decorator that might be what I need: https://numba.pydata.org/numba-doc/dev/extending/low-level.html#lower_cast

Secondly, there's an example of an ArrayCompatible in Numba's unit tests. The lowering of the __array__ method is here:

https://github.com/numba/numba/blob/13ece9b97e6f01f750e870347f231282325f60c3/numba/tests/pdlike_usecase.py#L180-L188