ENH: add a `segmented_reduce` method to ufuncs by ikrommyd · Pull Request #32243 · numpy/numpy

ikrommyd · 2026-08-10T18:52:36Z

PR summary

Taking over #25476.
Implements a new segmented_reduce method to ufuncs as it was discussed on that PR instead of enhancing the reduceat interface.
In its current state, this was mostly developed within the weekend by starting over by copying the whole reduceat code into a new segmented_reduce method, adding all the boilerplate for it and iterating with AI on making the necessary changes to it to support starts and stops offsets while considering the changes from Marten's PR. This is obviously not in a ready state and I have only skimmed through the code while iterating with AI. I just wanted a proof of concept over the weekend. I'm opening it as draft for visibility and because I wanted a CI run. I will be iterating over it in the next few days to get it in a ready to review state.

First time committer introduction

N/A

AI Disclosure

See above for the AI usage in the current state. Will update later.

cc @mhvk since the previous PR was yours

Signed-off-by: Iason Krommydas <iason.krom@gmail.com>

jorenham

sorry if premature, but I left some typing-related comments

jorenham · 2026-08-10T19:00:24Z


+add_newdoc('numpy._core', 'ufunc', ('segmented_reduce',
+    """
+    segmented_reduce($self, array, /, starts, stops=None, axis=0, dtype=None, out=None, **kwargs)


Can we make anything after stops=None keyword-only?

I will actually make stops required btw. I don't like the None case. Regarding keyword-only, there's no precedent for that in ufunc methods. No other ufunc methods have keyword-only args as far as I can see. We have to follow what reduce, reduceat and accumulate do here regarding their signature.

It's what the stubs enforce though, because I checked, and axis is rarely ever used as positional argument in practice for e.g. reduce.

there's no precedent for that in ufunc methods

Personally I don't think that consistency is worth repeating the mistakes of the past (even though I think consistency is pretty important).

Signed-off-by: Iason Krommydas <iason.krom@gmail.com>

github-actions · 2026-08-10T19:58:36Z

Diff from mypy_primer, showing the effect of this PR on type check results on a corpus of open source code:

optuna (https://github.com/optuna/optuna)
  ./optuna/study/_multi_objective.py:182: error: INTERNAL ERROR -- Please try using mypy master on GitHub:

mhvk

I have not looked in any detail yet, but think it is important that we end up in a state where there is only one function that deals with the internals of both reduceat and segmented_reduce, i.e., we should factor out the innards. Otherwise we're nearly guaranteed that the code will diverge.

Otherwise, only bikeshedding about names and keywords so far...

mhvk · 2026-08-10T23:58:15Z

+over segments (slices) of an axis given by their start and stop offsets::
+
+    >>> a = np.arange(8)
+    >>> np.add.segmented_reduce(a, starts=[0, 5], stops=[3, 8])


I'm not sure about the plural here, since a single value is in principle allowed.

My idea is that it comes from "start offsets" and "stop offsets" and offsets is always plural so I called them starts and stops. CCCL calls them "start_offsets_inandend_offsets_in`. I don't care much but I'm personally a fan of plural tbh.

My analogy is with slice, which has start and stop -- at some point I had a branch where for reduceat I let the user give a slice object with start and stop arrays. Also, things like np.linspace have start and stop which can be arrays.

But of course, for linspace, the expectation is that they are scalars, while here it is that you input multiple items with a single one an unusual possibility. Given that, I don't mind so much any more.

p.s. Actually, I think a single item should behave differently than a list of length 1: it should remove axis unless keepdims=True. If we do that, we can pass ufunc.reduce trivially through your new code... Relatedly, if the start and stop are >=2D arrays, one might want to insert multiple axes... For this PR, I suggest we insist on 1D arrays so that we do not prevent exploring use cases for different dimensions...

I was planning on doing ONLY 1D starts/stops anyways here. If somebody wants to expand that sure but indices in reduceat is also 1D only. I had only thought of segmentes along axis and nothing more than that. I'm haven't even thought about more than 1D 🤣

Excellent. The main thing is not to get starts and stops with something like the C-side equivalent of np.asanyarray(starts, ndmin=1), but rather explicitly check that they are 1D.

mhvk · 2026-08-10T23:59:44Z


+add_newdoc('numpy._core', 'ufunc', ('segmented_reduce',
+    """
+    segmented_reduce($self, array, /, starts, stops, axis=0, dtype=None, out=None, **kwargs)


I'm with @jorenham in thinking we should make everything starting with axis keyword only.

Generally I agree. My disagreement here is that this this feels like an improved version of reduceat to me and it feels nicer to maintain the way arguments are parsed consistent between them. Happy to do what the maintainers think of course. If I was designing the API from scratch I'd go keyword only for sure.

ikrommyd · 2026-08-11T15:33:58Z

I have not looked in any detail yet, but think it is important that we end up in a state where there is only one function that deals with the internals of both reduceat and segmented_reduce, i.e., we should factor out the innards. Otherwise we're nearly guaranteed that the code will diverge.

Otherwise, only bikeshedding about names and keywords so far...

Yes I started this by copying the reduceat code exactly as is. There's definitely a ton of duplicated code here overall. Reduceat also duplicates a lot of code from accumulate as well and there's a TODO in the code to fix that. I'd be willing to tackle this as well but I want to do this after the upgrade to handle reduction loops comes in.

mhvk · 2026-08-11T15:49:05Z

Yes, good to get the reduction, etc., loops in first -- they're a very nice improvement!

ikrommyd · 2026-08-11T15:52:11Z

Yes, good to get the reduction, etc., loops in first -- they're a very nice improvement!

Yeah treat this PR as "ready for design comments" basically 🤣. The generalization of the reduction loops to reduceat and accumulate (two PRs are open for this) will definitely come in first and this will be heavily rebased afterwards.

Shreyaparab28 · 2026-08-13T22:38:53Z

Great work

ENH: add a segmented_reduce method to ufuncs

f283a3a

Signed-off-by: Iason Krommydas <iason.krom@gmail.com>

github-actions Bot added the 01 - Enhancement label Aug 10, 2026

ikrommyd added this to the 2.6.0 Release milestone Aug 10, 2026

jorenham reviewed Aug 10, 2026

View reviewed changes

ENH: make stops a required argument of ufunc.segmented_reduce

0a5d4b6

Signed-off-by: Iason Krommydas <iason.krom@gmail.com>

This comment has been minimized.

Sign in to view

TST: only use offset dtypes that fit intp in segmented_reduce tests

a0f4c4b

Signed-off-by: Iason Krommydas <iason.krom@gmail.com>

mhvk reviewed Aug 11, 2026

View reviewed changes

Uh oh!

Uh oh!

Conversation

ikrommyd commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR summary

First time committer introduction

AI Disclosure

Uh oh!

jorenham left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ikrommyd Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

github-actions Bot commented Aug 10, 2026

Uh oh!

mhvk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ikrommyd Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ikrommyd commented Aug 11, 2026

Uh oh!

mhvk commented Aug 11, 2026

Uh oh!

ikrommyd commented Aug 11, 2026

Uh oh!

Shreyaparab28 commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ikrommyd commented Aug 10, 2026 •

edited

Loading

ikrommyd Aug 10, 2026 •

edited

Loading

ikrommyd Aug 11, 2026 •

edited

Loading