Skip to content

Add an axis argument to generalized ufuncs? #5197

@shoyer

Description

@shoyer

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?

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions