ENH: add mechanism to register reduction loops to ufuncs for multi-output reductions - #31816

CC @mhvk I would be very curious if you have some high-level thoughts on this?

@seberg - thanks for the ping!

@ikrommyd - this looks very nice indeed, and I like the general idea of having a new slot available to register reduction functions. Like you, I'm also pleased that in the end the changes are not very intrusive.

On the PR itself, one thing I think is important is to have something in numpy that exercises the new function; possibly, your minmax, maybe just for integers or floats, could serve that purpose? You'll have seen that there are several _core/src/umath/_*.c test functions, which you could add too.

Thanks @mhvk! Yeah I was planning to open a separate PR targetting this PR's branch that implements minmax in numpy but it wouldn't be much different than what the gist already has. I mean one just registers the reduction loops using the slot and PyUFunc_AddLoopFromSpec which is how normally one adds ufunc loops in the new-style array method API. I'm just working on something else too the last couple of days so I got caught with that. Let me know if you have any other suggestions about what you'd like to see. I was also planning to do this in an external extension like user dtypes or quaddtype to show how such extensions would interact with this API. I was also considering writing a short NEP about this too once the design is consolidated just for documentation purposes.

If Marten's OK with a minmax reduction landing as part of this PR, I'm fine with it too.

IMO a NEP isn't necessary to merge this, so it would be purely for documentation purposes. It would also add some process overhead that you might not want to deal with.

Yeah I wanted to keep the diff minimal for easier first look here on the design. We also need to touch reduceat and accumulate in a similar manner to reduce so that all 3 reduction-like ops work with this. Happy to add minmax + reduceat + accumulate changes here. I generally find smaller PRs easier to work with from the reviewers' perspective and then I would apply all the comments made here for reduce to accumulate and reduceat so what I originally had in mind was to have these and minmax as separate PRs but happy to adjust according to your liking 😃

All sounds good. I agree with the sense to stick to just reduce here, but include an example of some kind. I think your current minmax is perhaps more suitable as a simple test case, since it is not vectorized---I find it also really helpful to have those test examples... But just suggestions, do follow what works best for you, this is a nice contribution as is! One thing, as you were mentioning documentation, that definitely would need to be added, at least in the C API text.

Okay since I have not seen any opposing feedback on the design here, I plan to turn this into a more proper PR this week with testing and documentation. As Marten suggested, I will add the minmax toy example as a C test.
Then people can properly review and the feedback will be considered in the separate PRs that will touch accumulate and reduceat in a similar manner.

I've added documentation updates and I took my minmax and turned it into a C test by adding a double loop only + promoter for python scalars. I didn't want to add full dtype support as that will be done in the actual minmax implementation that will follow. The linting failure is unrelated and is on a file I'm not touching here.

IMO this is ready for the first review pass. There may be some doc displaying issues cause I didn't check the website preview but will do that later. Gonna do a rebase to main too.

Merge conflict is simple. It's just the new member in the array method struct from the loop caching PR. Will fix on the next rebase. Just don't wanna ruin the state right now.

I have implemented one more small get_multi_reduction_initials method on the array method to be able to set initial values to multi output reductions. Potentially we can deprecate the old get_reduction_initial and use this one which can definitely do the single output reductions too as Sebastian suggested considering earlier.

Whew that took a while to get done haha. I feel like I've addressed all your review comments @ngoldbaum!

@ikrommyd OK, I think this is ready now. You'll need to clear the merge conflict and then I think I'll go ahead and merge.

While you're doing that, I also think this test would add some coverage for the object forward loop:

def test_object_forward(self):
    a = np.array([3, -7, 12], dtype=object)
    b = np.array([1, 9, -2], dtype=object)
    got_min, got_max = mm(a, b)
    assert got_min.tolist() == [1, -7, -2]
    assert got_max.tolist() == [3, 9, 12]

@ngoldbaum - rebased, fixed the conflict and added the test you suggested.

OK, no one else weighed in since I suggested merging this yesterday, so in it goes!

If anyone would like to register any comments or suggestions, just ping Iason and he'll get to it on a second pass.