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], [5], [6, 7, 8, 9]]))
array([ 3., 0., 7., 5., 30.])
because input is an iterable Awkward Array yielding Awkward Arrays x and x is an iterable Awkward Array yielding numbers y, and 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], [5], [6, 7, 8, 9]]))
As before, 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 float or bool, then it should be registered as convertible to both float and 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?