Feature request: axis and keepdims for gufuncs (and use it in np.median) #8810

mhvk · 2017-03-22T15:05:44Z

EDITS following different comments:

(Feature request: axis and keepdims for gufuncs (and use it in np.median) #8810 (comment)): include keepdims. See also previous discussion in Add an axis argument to generalized ufuncs? #5197 and on https://mail.scipy.org/pipermail/numpy-discussion/2014-October/071455.html (which independently suggested more or less the same syntax).
see ENH: Implement axes keyword argument for gufuncs. #8819 for the implementation of axes; adjusted implementation notes accordingly.

Rationale

Following the suggestion that, ideally, np.median and similar functions would be gufuncs so that they could be overridden with __array_ufunc__ (#8247), it was realised that the main show-stopper is the absence of an axis argument for gufuncs (also annoying in the context of the new all_equal gufuncs (#8528). Furthermore, for functions like argmax, etc., a keepdims would be very useful (#8710).

Specification

In practice, the new axis argument would use underneath a more general axes argument, which is a list of tuples of axes for each of the core dimensions. For normal ufunc, axes obviously has to be empty (or absent); for gufunc, defaults would be -1, -2, etc., as needed. For the output, the tuple would be extended by one (or possibly more) from that implied by the signature if keepdims is set . Short-cuts may be possible, e.g., for gufuncs with only one index, a single number (or perhaps even None for ravelling all dimensions); one would distinguish the general case of axes with the simple case of axis.

Implementation

Looking at the reference, it would seem that the C-api would not need to change, but the routine actually calling the underlying iteration machinery in umath/ufunc_object.c would need to do appropriate transposes of the input and output arrays. Going through the call sequence to see places where change would be necessary (not marked with X):

ufunc_generic_call (all OK)
PyUFunc_CheckOverride (OK, axis should already be in kwds)
PyUFunc_GenericFunction: signal that non-empty axis should error.
PyUFunc_GeneralizedFunction: do a remap of axes for the checks and before passing on to the iterator;
get_ufunc_arguments: interpret axes argument (new return argument); error if not needed. ENH: Implement axes keyword argument for gufuncs. #8819
Implement the simple axis one for just a single axis. ENH: Implement axis for generalized ufuncs. #11018
repeat above for keep_dims; ENH: Implement axis for generalized ufuncs. #11018

Excluded

In principle, it might be nice to allow axis to refer to multiple axes or all of them (None), but this is substantially more complicated, and it is not obvious it isn't better to let wrapping code deal with this. Indeed, this is done for np.median, which ravels any requested axes as needed before passing the array on to partition (which can handle only a single axis).

The text was updated successfully, but these errors were encountered:

eric-wieser · 2017-03-22T15:16:04Z

There's some extended discussion about this on the mailing list. (and in #5197)

I think it might be worth special-casing signatures of the form (n)->() first, and additionally adding a keepdims argument (#8710).

mhvk · 2017-03-22T15:26:54Z

Ah, yes, I should have checked for duplication; maybe this issue can be specifically about implementation. In that respect, good point about keepdims -- I've added it in the description.

charris · 2017-03-22T15:29:29Z

I don't think gufuncs are really a good fit for an axis argument as they are currently conceived. It might be good to make a list of the functions we would like to implement as ufuncs to see what would be the best solution.

eric-wieser · 2017-03-22T15:29:51Z

Thinking more carefully, any reduction of the form (m,n,...)->() or even (m,n,...),(o,p,...)->() (where ... are matching lengths) has an unambiguous meaning of keepdims, but beyond the (n)->() case, we need to think very carefully about the interpretation of the axis argument

mhvk · 2017-03-22T15:39:16Z

@charris - I think it will actually be quite trivial to implement (but realise that I may have to prove that...)

shoyer · 2017-03-22T16:06:46Z

I don't think adding axis to gufuncs helps that much for median, because the default behavior for median is to reduce over all axes (also consider a single axis vs. tuples of various lengths).

We could certainly add another override for functions with the standard logic for handling axis on aggregation functions, but it should probably be more specialized that __array_ufunc__.

mhvk · 2017-03-22T16:12:42Z

@shoyer - maybe median is not the perfect example, but it would still help: it would need a separate front-end that, for axis=None or multiple axes, would ravel the array appropriately. But that is also all it needs to do before passing it on to the gufunc implementation, and with just that, subclasses could still override the actual calculation, as long as they allow reshaping (which I hope would be the majority).

p.s. Just to be clear: my goal is not to make every numerical function fit in the gufunc mold, but to ensure the actual numerical aspect of all those functions can be overridden by duck-type arrays.

charris · 2017-03-22T16:44:36Z

I'm not so much interested in the difficulty as the usefulness. Would it be correct to say that the main interest in adding an axis is to reuse the override functionality?

mhvk · 2017-03-22T17:02:31Z

@charris - It is a combination of liking the idea of being able to override numerical functions and direct usefulness. For the latter, I am thinking in particular of #8528 (all_equal), which to me showed the weakness of the present scheme, and the realisation in #8710 that argsort and co would match the gufunc behaviour well if only it supported axis and keepdims.

mhvk · 2017-03-23T21:08:35Z

See #8819 for a trial implementation of adding an axis keyword. Since the iterator supports fairly arbitrary axis, it was mostly a matter of adding a layer of indirection in the code preparing arrays.

Added notes:

A few quick timing tests show a 50% increase in time if one passes in axes for small arrays; for large ones, it depends only on memory organisation of the arrays.
Haven't really thought about keepdims yet.

mhvk · 2018-07-03T20:53:16Z

Looking at this issue again, I realized it can actually be closed: gufuncs now have axis and keepdims, so functions like partition could be rewritten in terms of gufuncs if so desired. Possible raveling of multiple axes would still need to be done outside of the gufunc.

mhvk added the 23 - Wish List label Mar 22, 2017

mhvk mentioned this issue Mar 22, 2017

ENH: Adding logical gufuncs, ie all_equal #8528

Closed

mhvk changed the title ~~Feature request: axis argument for gufuncs (and use it in np.median)~~ Feature request: axis and keepdims for gufuncs (and use it in np.median) Mar 22, 2017

mhvk mentioned this issue Mar 22, 2017

Add an axis argument to generalized ufuncs? #5197

Closed

mhvk mentioned this issue Mar 22, 2017

Make median a member of the class numpy.ndarray #8797

Closed

mhvk mentioned this issue Mar 29, 2017

ENH: Implement axes keyword argument for gufuncs. #8819

Merged

mhvk closed this as completed Jul 3, 2018

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature request: axis and keepdims for gufuncs (and use it in np.median) #8810

Feature request: axis and keepdims for gufuncs (and use it in np.median) #8810

mhvk commented Mar 22, 2017 •

edited

Loading

eric-wieser commented Mar 22, 2017 •

edited

Loading

mhvk commented Mar 22, 2017

charris commented Mar 22, 2017

eric-wieser commented Mar 22, 2017 •

edited

Loading

mhvk commented Mar 22, 2017

shoyer commented Mar 22, 2017

mhvk commented Mar 22, 2017 •

edited

Loading

charris commented Mar 22, 2017

mhvk commented Mar 22, 2017

mhvk commented Mar 23, 2017 •

edited

Loading

mhvk commented Jul 3, 2018

Feature request: axis and keepdims for gufuncs (and use it in np.median) #8810

Feature request: axis and keepdims for gufuncs (and use it in np.median) #8810

Comments

mhvk commented Mar 22, 2017 • edited Loading

Rationale

Specification

Implementation

Excluded

eric-wieser commented Mar 22, 2017 • edited Loading

mhvk commented Mar 22, 2017

charris commented Mar 22, 2017

eric-wieser commented Mar 22, 2017 • edited Loading

mhvk commented Mar 22, 2017

shoyer commented Mar 22, 2017

mhvk commented Mar 22, 2017 • edited Loading

charris commented Mar 22, 2017

mhvk commented Mar 22, 2017

mhvk commented Mar 23, 2017 • edited Loading

mhvk commented Jul 3, 2018

mhvk commented Mar 22, 2017 •

edited

Loading

eric-wieser commented Mar 22, 2017 •

edited

Loading

eric-wieser commented Mar 22, 2017 •

edited

Loading

mhvk commented Mar 22, 2017 •

edited

Loading

mhvk commented Mar 23, 2017 •

edited

Loading