@shoyer (commenting in main thread, since a good point): For reduce operations the usage of axis is a bit different: one can tell to go over multiple rather than a single axis. In contrast, here we want to define where the core dimensions of the gufunc are to be found.

Now a counterargument is that np.sum (and any reduce operation) could be seen as a gufunc with signature (i)->(). In that respect, what the tuple does is, effectively, combining those axes (with None combining all).

But perhaps that argues that the items inside the list should also be lists, so that we keep open the option of using tuples to define multiple operand axes as a single core axis. I'll do that...

But perhaps that argues that the items inside the list should also be lists, so that we keep open the option of using tuples to define multiple operand axes as a single core axis. I'll do that...

How about using sets for that meaning? None of the reducer functions care about the ordering of their operand axes, right?

Alternatively: I'd be tempted to allow the reducers to be defined as (i...), which explicitly allows this dimensionality reduction (and allows keepdims to restore them). Then the caller can just pass a tuple of indices to this as they would (i,j,k)

Right now, do gufuncs incur a copy if their non-core dimensions are not contiguous?

@eric-wieser - A set would be logical, but I think @shoyer's point was in part that for np.sum, etc., a tuple is already used.

@eric-wieser - I'm not sure I understand what you mean by the "reducer". Would this be in the signature?

For the non-core dimensions, an iterator is used, so no copy is made (strides are passed to the inner loop, so no copy is made there either).

@eric-wieser - I'm not sure I understand what you mean by the "reducer". Would this be in the signature?

I mean the functions with a signature ending in ->(), which are also the only ones that should allow keepdims IMO. I guess reduce is a bad word, because argmin is not a reduce operation.

Yes, (i...) would be the literal signature. This would need special handling in the inner loop, since that would now need to iterate over a variable number of strides.

I think I'd prefer to keep discussion of the signature out of this PR -- it is complicated enough as it is!

I also think you're right about keepdims making sense only for an output with no core dimensions (and probably only when all inputs have the same core dimensions).

I'm also wondering if, rather than special-case so much, it is better to go with @njsmith's suggestion of having axis for the simple case of just one core dimension, and axes always requiring full details.

I think I'd prefer to keep discussion of the signature out of this PR -

My point here is that defining the signature in this way would let us use tuples for the inner axis argume without ambiguity

• [(1, 2), 3] for (m,n)->(n), as previously proposed
• (1, 2) for (m,n)->(), as previously proposed
• (1, 2) for (m...)->(), which would implement sum(axis=(1,2))

Hmm, I had thought that, really, the np.sum one corresponded to treating several axes as a single one. And that, thus, it would be allowed for any input core dimension to specify multiple axes. But at some level that can already be done by an appropriate reshape inside the actual call (though that argument also holds for np.sum...). I'll think a bit more.

it would be allowed for any input core dimension to specify multiple axes

I'm struggling to think of a case when this is useful and not just cryptic.

can already be done by an appropriate reshape inside the actual call (though that argument also holds for np.sum...)

Except we don't want to do that, (and I think currently don't?), because reshape would cause a copy. We need an nd-iter at that point, and the ufunc machinery starts to fall apart

But .reduce can go over multiple axes. Looking at the reduction.c, I found this gem:

 * TODO FIXME: if you squint, this is essentially an second independent
 * implementation of generalized ufuncs with signature (i)->(), plus a few
 * extra bells and whistles. (Indeed, as far as I can tell, it was originally
 * split out to support a fancy version of count_nonzero... which is not
 * actually a reduction function at all, it's just a (i)->() function!) So
 * probably these two implementation should be merged into one. (In fact it
 * would be quite nice to support axis= and keepdims etc. for arbitrary
 * generalized ufuncs!)

(no surprise now that @njsmith has been bemused by our discussion...)

@jaimefrio - thanks for the comments! Your way to fill remap_axis is much better; I now keep a flag for whether remapping is necessary at all, so I could just use it verbatim. (Well, with better names, also following your suggestions.)

I also agree with you that this better go past the mailing list! To help that, I've started writing some test cases (now included) and will add documentation. I'm still pondering whether to try to get keepdims working before bothering people.

On the syntax: I thought a single axis still work for a (i),(i)->() or (i),(i)->(i) signature. For instance, I could imagine using the new all_equal on vectors, where the two inputs would not necessarily have xyz in the last dimension.

To all those still interested: I went back to a list of tuples, which I find much clearer (perhaps it should even be a tuple of tuples), and I ended up not seeing how we could ever sensibly support turning multiple axes into one anyway.

Responses to @mhvk's first and latest posts are below:

In general, axis should be a list with length equal to either the number of inputs or the total number of operands (inputs & outputs).

Instead, I might say: axis should be a list with length equal to the number of inputs with a core dimension, or the total number of inputs and outputs with core dimensions.

The following shortcuts are supported (to be confirmed with mailing list): an entry in the (outer) list of None is taken to imply the default

I don't see much point in this, unless you have a large number of core dimensions. Writing -1 is more explicit and requires less text, too.

Alternatively, we could use None as a sentinel for "all dimensions", which is harder to write (i.e., tuple(range(-array.ndim, 0, 1)), which took me some experimentation to figure out).

an integer is OK if the operand has only a single core dimension.

Yes, this makes sense.

Furthermore, if the signature has only one core dimension, then a simple integer can replace the outer list.

By one core dimension, do you mean one core dimension on exactly one argument?

On the syntax: I thought a single axis still work for a (i),(i)->() or (i),(i)->(i) signature.

Do you still think this works? I'm not sure how. "Broadcasting" an integer axis over multiple arguments seems like a slippery road to go down.

I think we have discussed this before, but I wanted to bring up again that one use of an axis argument is to control dispatching to gufunc kernels, even for higher dimensional inputs.

e.g., for matmul:

• np.matmul(x, y, axis=[(0,), (0,)] could dispatch to the kernel (n),(m)->()
• np.matmul(x, y, axis=[(0, 1), (0,)] could dispatch to the kernel (n,m),(m)->(n)
• np.matmul(x, y, axis=[(0,), (0, 1)] could dispatch to the kernel (n),(n,m)->(m)

We don't need to support this yet, but we should be sure we don't do anything that precludes it.

For reference, the Python parser I wrote for handling gufunc axis argument can be found here: http://nbviewer.jupyter.org/gist/shoyer/7740d32850084261d870

I'm not sure how helpful this is, but probably it's still a good idea to prototype the logic for axis canonicalization in Python.

Instead, I might say: axis should be a list with length equal to the number of inputs with a core dimension, or the total number of inputs and outputs with core dimensions.

I'm not sure - I think it's easier to understand when there's a 1-to-1 mapping between arguments and axes - even if it comes at the cost of having to specify () for arguments with no core dimensions

I don't see much point in this, unless you have a large number of core dimensions. Writing -1 is more explicit and requires less text, too. [...] Alternatively, we could use None as a sentinel for "all dimension".

+1 from me on this. For now, let's just reserve None for that use, rather than actually implementing it.

On the syntax: I thought a single axis still work for a (i),(i)->() or (i),(i)->(i) signature.

The former would be something like all_equal, and the latter np.cross? I guess that seems reasonable.

I'm not sure how helpful this is, but probably it's still a good idea to prototype the logic for axis canonicalization in Python.

Can we just do the canonicalization in _internal.py? Presumably we'll want to pass a canonicalized axis to __array_ufunc__ anyway, so need to go via the full list of tuples

@shoyer, @eric-wieser - OK, I agree, let's reserve None for possible different use. I now removed that from the code. I agree with Eric that an empty tuple for an operand with no core dimensions is cleaner (note that in many cases, this will be the output, which, following @jaimefrio's suggestion, can now be omitted).

@shoyer - did the examples from @eric-wieser and myself convince you it makes sense to keep the single-integer case as applying to any single core dimension?

I very much like the idea of being able to use the axis as a way to select the signature in, e.g., matmul, even though that would quite obviously need a much more major intervention (since inside PyUFunc_GeneralizedFunction, I'd seem stuck with a single signature ufunc). But I don't think the current API prevents this from happening at a later stage.

I do agree that, in principle, one could move the translation from integer to single-element tuple (or set of single-element tuples) outside, although this does mean that ufunc_override.c would have to become rather uncomfortably aware of ufunc internals. Anyway, I would propose to leave this for later... (but do have a look at #8247 -- which (re)implements __array_ufunc__ -- if you want to think about the interaction).

@shoyer - did the examples from @eric-wieser and myself convince you it makes sense to keep the single-integer case as applying to any single core dimension?

I think I missed those examples -- can you point to them again?

One option that I believe @njsmith and I discussed previously is is to have two separate keyword arguments:

• axes for providing the full axes specification in long form
• axis for specifying an integer axis.

I would actually lean toward only making axes part of the gufunc specification, and making handling of an axis argument function specific, e.g., if we wanted the core of numpy.cumsum to use a gufunc, we could write:

def cumsum(array, axis):
    axes = [(axis,), (axis,)]
    # cumsum_impl is a gufunc with signature (n)->(n)
    return cumsum_impl(array, axes=axes)

This allows for the simplifying logic of an axis argument when it makes sense, but keeps the implementation specific to each function that uses it.

The only slightly confusing part about this is that implementors of __array_ufunc__ would need to check for cumsum_impl instead of cumsum if they override it's functionality.

@shoyer - the case for allowing axis=i for any set of single core-dimension signatures, is that for the examples we could come up with, it seemed obvious what was the intent:

• (i),()->(): we all agree
• (i),(i)->(): for something like all_equal or inner1d it would make sense certainly if one knows one is passing in data with the same structure (e.g., vector arrays where the xyz dimension is 0 rather than -1).
• (i),(i)->(i): same, e.g., a cross-product of the above vectors.

I can see the appeal of axis vs axes, though, with each fairly strictly defined. I mostly ended up not (yet) doing it as it seemed not quite necessary. But if, overall, people feel this whole PR is going in the right direction, and it would be better with axis and axes, I'm happy to implement it. (of course, PRs to my PR are even more welcome...)

I do wonder a bit about keepdims. Earlier it was suggested to only allow it for the case of a single core dimension, but that would remove, e.g., the possibility to let a matrix determinant ufunc with signature (i,i)->() return an array with shape ending in (1,1). One might just more generally add 1 as necessary to match the longest input core dimension.

LGTM, though I didn't check super carefully for little bugs.

If I understand correctly, a future plan may be to allow collapsing multiple axes to one axis like reduce does. The syntax might allow nested tuples like [((-4,-3), -1), (-1,), (-2,)] for signature (i,j),(j)->(i) where for the 0th op axes -4 and -3 were collapsed to a single axis.

That's starting to get a bit cluttered, so I was thinking about other syntaxes. What about supporting a dictionary like axes={0: ((-4,-3), -1), 2: (-2,)} where the key is the op index and the value is the axes? That gives a way to only specify some op axes and leaves the others as default, like you attempted with the None placeholder, now gone. And indstead of having to mentally count out which op a tuple corresponds to, you can just read off the key.

@ahaldane - I quite like the dict idea, in particular that it is so much more obvious that one can override the axes for just some operands. On the other hand, it feels like it deviates a bit from "standard" practice, in that maybe more typically extra attributes to what in essence is a tuple of arguments are also given as iterables. But I'm not sure this "standard" practice really exists anywhere but in my mind.

I also feel having keys that are just numbers is slightly odd, but perhaps a related alternative would resolve this: one could look in kwargs for any keyword arguments starting with axes (i.e., axes0, axes1; still plural, since it can be more than one) and pass those on. In terms of implementation, this would likely end up quite similar to your suggestion, since it would be easiest to just create a dict with just what you suggest in the parser.

@eric-wieser, what do you think?

@mhvk I quite like the keyword arg idea. I also found some other alternatives in this thread.

Here's an example of what the different options we've mentioned look like:

np.matmul(a, b, out=c, axes=[((-4,-3), -1), (-1,), (-2,)])

np.matmul(a, b, out=c, axes={0: ((-4,-3), -1), 2: (-2,)})

np.matmul(a, b, out=c, axes0=((-4,-3), -1), axes2=(-2,))

np.matmul(a, b, out=c, axes0=((-4,-3), -1), axes_out0=(-2,))

np.matmul(a, b, out=c, axes=[{"m": (-4, -3), "n": -1}, {"n": -1}, {"m": -2}].

I think my favorites right now are the axes0= styles.

The nested tuple syntax feels most natural to me, though I agree that it can get cumbersome for complex expressions.

I would be OK with optionally replacing tuples at either the argument or axis level with dictionaries, keyed by either argument number or axis name, e.g., any of

• axes={0: ((-4,-3), -1), 2: (-2,)})
• axes=[{"m": (-4, -3), "n": -1}, {"n": -1}, {"m": -2}]
• axes={0: {"m": (-4, -3), "n": -1}, 2: {"m": -2}}

That said, I'm skeptical that we actually need this. If the main argument is readability, then dictionary for axis names is probably the most valuable (e.g., {"m": (-4, -3), "n": -1}).

I don't like the axis0= styles because they rely on programmatically generated argument names.

That sounds fair, let's support the list-of-tuples format in any case. As far as I can see, everyone is happy with this PR and it is ready to merge.

So how about we merge this PR after 1.14.1 gets released in the next few days? We can leave any additional syntaxes or options like 'keepdims' for discussion in future PRs, aiming for 1.15.

Sounds all OK to me.

I think my favorites right now are the axes0= styles.

Would they scale well to an arbitrary amount if input arguments? It seems einsum can easily call a gufunc version of a fused-multiply-add if arbitrary amounts of arguments were kept.

My favourite is the int-indexed dict style. A perfect balance of verbosity and explicitness, IMO.

axes={0: ((-4,-3), -1), 2: (-2,)}

Now that 1.14.1 is in, let's go ahead and merge this one. Then we can debate alternate syntaxes.

I'll give this one another read-through myself, and then I plan to merge in half a day. Sounds good?

If anyone wants more time to check anything, or if anyone who did the more thorough reviews above wants to give the final say, please speak up before then.