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 <= arr.shape: acc = arr.copy() for v in iter(arr[1:]): acc = func(acc, v) else: acc = arr.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.