-
-
Notifications
You must be signed in to change notification settings - Fork 12k
Description
I would like to write generalized ufuncs, using numba, to create fast functions such as nanmean (signature '(n)->()') or rolling_mean (signature '(n),()->(n)') that take the axis along which to aggregate as a keyword argument, e.g., nanmean(x, axis=0) or rolling_mean(x, window=5, axis=0).
Of course, I have the option of writing my own wrapper function to do something like x.swapaxes(axis, -1) before passing the array off to the gufunc. This is not so terrible.
However, it occurs to me that this sort of dimension reordering might be handled more cleanly (and efficiently) within the gufunc machinery itself, somewhat analogously to numpy.ufunc.reduce, except the axis argument (if supplied) would need to be an integer or tuple with length equal to the number of core dimensions on the first gufunc argument.
Thoughts?