Using Numba’s extension API, is it possible to extend all functions and builtin operators that expect type Y with type X by inserting an X → Y conversion function?
Here’s the use-case: hundreds of functions taking NumPy arrays have been defined in core Numba. I have extended an array type that is sometimes convertible to NumPy arrays (I can raise exceptions in the conversion function if it is not). I would like to implement all of these functions and any that may be added in the future by registering an implicit conversion, similar to Scala’s implicit conversions.
For example, I can implement sums over contents of Awkward Arrays like this:
>>> import numpy as np >>> import numba as nb >>> import awkward1 as ak >>> @nb.njit ... def f(input): ... output = np.zeros(len(input), np.float64) ... for i, x in enumerate(input): ... for y in x: ... output[i] += y ... return output >>> f(ak.Array([[0, 1, 2], , [3, 4], , [6, 7, 8, 9]])) array([ 3., 0., 7., 5., 30.])
input is an iterable Awkward Array yielding Awkward Arrays
x is an iterable Awkward Array yielding numbers
output[i] += y knows how to add numbers to items of an array.
But suppose I want to write it like
>>> @nb.njit ... def f(input): ... output = np.zeros(len(input), np.float64) ... for i, x in enumerate(input): ... output[i] = np.sum(x) ... return output ... >>> f(ak.Array([[0, 1, 2], , [3, 4], , [6, 7, 8, 9]]))
x is an Awkward Array. I can lower a conversion function
to_numpy that converts the Awkward Array into a (lowered) NumPy array or raise an exception trying. However, I need
- Numba’s typing pass to recognize all functions, not just
np.sum, that take a NumPy array as also being open to a signature with an Awkward Array in place of the NumPy array, and
- Numba’s lowering pass to insert my conversion function in those places.
Is there anything I can do about this? The use-case described above is scikit-hep/awkward-1.0#509, but this would also enable scikit-hep/awkward-1.0#174, implementing union-typed Awkward Arrays, because that also has to overload an open-ended set of functions. (If the union array can have values of type
bool, then it should be registered as convertible to both
bool, but if a value wants to resolve to
bool and the particular datum happens to be
float, the conversion function would raise an exception. That implements dynamic type-checking for a fixed set of types.)
The convertible-to-array case (issue #509) would be simpler to implement and is a more immediate need than union arrays (issue #174), but they both depend on this capability.
Is this already possible in Numba or would something have to change?