Multiple Returns for an Overloaded Attribute Function

Numba Support: How do I do ...?
1

Hi all,

I’m implementing ufunc.reduce with axis arguments based on a conversation with stuartarchibald in Gitter earlier and hoping to extend it to the other four ufunc methods later. Example code:

@register_jitable
def flat_reduce_func(arr, func):
    arr_flat=arr.ravel()
    acc = 0
    for v in np.nditer(arr_flat):
        acc = func(acc, v.item())
    return acc

@register_jitable
def axis_reduce_func(arr, axis, func):
    if arr.shape[0] <= arr.shape[1]:
        acc = arr[0].copy()
        for v in iter(arr[1:]):
            acc = func(acc, v)
    else: 
        acc = arr[0].copy()
        for idx, v in enumerate(arr.T):
            put = 0
            for v_2 in iter(v):
                put = func(put, v_2.item())
            acc[idx] = put
    return acc  

@overload_attribute(types.Function, 'reduce')
def ol_np_ufunc_reduce(self):
    if isinstance(self.typing_key, np.ufunc):
        func = self.typing_key
        @njit
        def reduce_func(arr, axis=0):
            if axis is None:
                return flat_reduce_func(arr, func)
            else:
                return axis_reduce_func(arr, axis, func)
            return reduce_func
          
        def impl(self):
            return reduce_func
        return impl

This works well, but I’m having type unification issues with the returns for reduce_func, as it is currently returning either a single value or an array. If this was a function, I’d deal with that by overloading it, but I’m actually not sure what to do here as I can’t find any examples of overload_attribute with optional arguments in the main repo. Any ideas?

Separately, ufunc.reduce is a method, really, so I imagine that there’s probably a way to overload_method this. Again, there’s not much adaptable literature out there.

Thanks!