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. I tried to make the code changes as minimal as possible and take most of the implementation from reduceat from helpers to reduce duplication.

AI Disclosure

I iterated with AI on most code outside the actual ufunc_object.c implementation like tests, typing, and boilerplate. The core implementation was handled more manually.

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).

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.

ikrommyd · 2026-08-17T21:41:43Z

This is honestly mostly ready from my side. It just needs code-deduplication with reduceat which is easy when I look at the diff between the two. However, this would conflict with the work done on multi-output reductions for reduceat and accumulate and I think those should be merged first cause then we'd want to also make multi-output reductions work for segmented_reduce. But anyone is free to take a look as the core logic is down here now.

ikrommyd · 2026-09-24T03:49:28Z

I'm keeping this a bit frozen until multi-output reductions land in main for reduceat but I had a question in the meantime. Do you think segmented reduce should support multi-output reductions? Sure completeness is an argument but I thought maybe not yet to keep the code of segmented reduce simpler? I don't know. I doubt anyone eagerly wants a segmented minmax for example.

mhvk · 2026-09-24T13:11:49Z

@ikrommyd - I think deduplication with reduceat would be the highest priority. Eventually, one would also like to deduplicate with reduce itself, and it would seem that at that point one would get the multi-output work automatically. But it seems reasonable (and probably in fact best) to do it as two steps.

ikrommyd · 2026-09-24T14:16:23Z

@ikrommyd - I think deduplication with reduceat would be the highest priority. Eventually, one would also like to deduplicate with reduce itself, and it would seem that at that point one would get the multi-output work automatically. But it seems reasonable (and probably in fact best) to do it as two steps.

yeah part of why I keep this frozen for now is because reduceat is changing in #32212. So I want to restart with that PR (once merged) as a basis.

… values Signed-off-by: Iason Krommydas <[email protected]>

Signed-off-by: Iason Krommydas <[email protected]>

…uce` Signed-off-by: Iason Krommydas <[email protected]>

…educe Signed-off-by: Iason Krommydas <[email protected]>

Signed-off-by: Iason Krommydas <[email protected]>

…mented_reduce` Signed-off-by: Iason Krommydas <[email protected]>

ikrommyd · 2026-11-09T11:36:20Z

@jorenham should I implement the typing here or want me to leave it for a future PR or for you? it's currently in the diff as copy-pasted from reduceat with the names changed.

jorenham · 2026-11-09T11:38:29Z

@jorenham should I implement the typing here or want me to leave it for a future PR or for you? it's currently in the diff as copy-pasted from reduceat with the names changed.

Copy-pasting is probably what I'd do myself, so I see no reason to defer that. I'll take a closer look at the diff once this is ready for review.

ikrommyd · 2026-11-09T11:41:47Z

Yeah I just did a massive rebase as most of the reduceat core has been changed to add the multi-output support. This is now in sync with main's reduceat implementation and also supports multi-output loops out of the box. I want to do a pass myself to simplify more if possible and an AI review but I guess the typing shouldn't change so if you're ever bored you can do a pass on that already I guess.

…ented_reduce` Signed-off-by: Iason Krommydas <[email protected]>

…uceat_impl` Signed-off-by: Iason Krommydas <[email protected]>

Signed-off-by: Iason Krommydas <[email protected]>

jorenham · 2026-11-11T14:30:59Z

+        axis: SupportsIndex = 0,
+        dtype: DTypeLike | None = None,
+        out: EllipsisType | None = None,
+        initial: Incomplete = ...,


Below in the __init__.pyi changes for ufunc you added a comment saying that initial is keyword-only, but right now it can also be passed as positional arg here.

Could we narrow Incomplete to something like _ScalarLike_co or ArrayLike? Or, if it can accept anything, then we should use object instead of Incomplete, because Incomplete is basically intended as a # TODO: annotate me.

Is there a default value we could use here instead of the = ...? Or is it more like an if "initial" in kwargs: ... kinda situation?

jorenham · 2026-11-11T14:38:19Z

IIRC you said that the segmented_reduce signature in these ufunc stubs is pretty much equivalent to that of reduceat (modulo the extra params that aren't tested here). Because if that's the case, we don't need to also copy-paste their type-tests here, and we could just write a simple comment like we do elsewhere, e.g.:
https://github.com/ikrommyd/numpy/blob/36312d6d1611ef6a3878d4bbf4e849f476b8d4b9/numpy/typing/tests/data/reveal/ufuncs.pyi#L3928

jorenham · 2026-11-11T14:39:27Z

+assert_type(np.frompyfunc(func21, n2, n1).segmented_reduce(AR_f4_2d, AR_i8, AR_i8), npt.Array2D[np.object_[int]])
+assert_type(np.frompyfunc(func21, n2, n1).segmented_reduce([f8], AR_i8, AR_i8), npt.Array1D[np.object_[int]])
+assert_type(np.frompyfunc(func21, n2, n1).segmented_reduce([[f8]], AR_i8, AR_i8), npt.Array2D[np.object_[int]])
+assert_type(np.frompyfunc(func21, n2, n1).segmented_reduce(AR_f8, AR_i8, AR_i8, out=AR_f8), npt.NDArray[np.float64])
+assert_type(np.frompyfunc(func21, n2, n1).segmented_reduce([AR_f8], AR_i8, AR_i8), npt.NDArray[np.object_[int]])


If this is also the same as reduceat here, then a comment would suffice

jorenham · 2026-11-11T14:41:50Z

+        starts: _ArrayLikeInt_co,
+        stops: _ArrayLikeInt_co,


_ArrayLikeInt_co also accepts boolean array-likes, which I'm guessing wouldn't work at runtime. So let's use _ArrayLikeInt here instead (without the _co), like you already did above in _core.umath. It's not yet imported though, so you'll have to import it from ._array_like.

jorenham · 2026-11-11T14:45:31Z

+        # `initial` is passed through `kwargs` (like it is for `reduce`)
+        **kwargs: Incomplete,


No need to repeat the mistakes of the past: it's easy enough to instead write something like

Suggested change

# `initial` is passed through `kwargs` (like it is for `reduce`)

**kwargs: Incomplete,

*,

initial: _ScalarLike_co,

jorenham · 2026-11-11T14:45:56Z

+        starts: _ArrayLikeInt_co,
+        stops: _ArrayLikeInt_co,


Suggested change

starts: _ArrayLikeInt_co,

stops: _ArrayLikeInt_co,

starts: _ArrayLikeInt,

stops: _ArrayLikeInt,

jorenham · 2026-11-11T14:52:14Z

+    segmented_reduce($self, array, /, starts, stops, axis=0, dtype=None, out=None, **kwargs)
+    --
+
+    segmented_reduce(array, starts, stops, axis=0, dtype=None, out=None, initial=<no value>)


should we also mark the pos-only and kw-only params here?

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

This comment has been minimized.

Sign in to view

mhvk reviewed Aug 11, 2026

View reviewed changes

ikrommyd added the component: numpy.ufunc label Oct 7, 2026

ikrommyd added 7 commits November 9, 2026 00:53

MNT: generalize the reduceat segment loop to stop offsets and initial…

cc92dfd

… values Signed-off-by: Iason Krommydas <[email protected]>

ENH: add a segmented_reduce method to ufuncs

db07f2c

Signed-off-by: Iason Krommydas <[email protected]>

ENH: support ufuncs with more than one output in `ufunc.segmented_red…

6fe902c

…uce` Signed-off-by: Iason Krommydas <[email protected]>

MNT: share the per-output initial packing of reduce and segmented_r…

d99eda4

…educe Signed-off-by: Iason Krommydas <[email protected]>

BUG: do not broadcast to a larger out in ufunc.reduceat

26619d3

Signed-off-by: Iason Krommydas <[email protected]>

TYP: type ufunc.segmented_reduce like ufunc.reduceat

6176002

Signed-off-by: Iason Krommydas <[email protected]>

BUG: get the initial value of empty segments separately in `ufunc.seg…

b9418ab

…mented_reduce` Signed-off-by: Iason Krommydas <[email protected]>

ikrommyd force-pushed the segmented-reduce branch from 595511d to b9418ab Compare November 9, 2026 11:30

ikrommyd added 2 commits November 9, 2026 17:57

MNT: mention initial=None in the empty segment error of `ufunc.segm…

8fc8fac

…ented_reduce` Signed-off-by: Iason Krommydas <[email protected]>

MNT: split PyUFunc_Reduceat into thin wrappers around a shared `red…

06b8001

…uceat_impl` Signed-off-by: Iason Krommydas <[email protected]>

ikrommyd marked this pull request as ready for review November 9, 2026 12:44

ikrommyd requested a review from ngoldbaum November 10, 2026 15:14

MNT: some styling changes

36312d6

Signed-off-by: Iason Krommydas <[email protected]>

jorenham self-requested a review November 11, 2026 14:26

jorenham reviewed Nov 11, 2026

View reviewed changes

		# `initial` is passed through `kwargs` (like it is for `reduce`)
		**kwargs: Incomplete,

Conversation

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

Uh oh!

PR summary

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.

This comment has been minimized.

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!

ikrommyd commented Aug 17, 2026

Uh oh!

ikrommyd commented Sep 24, 2026

Uh oh!

mhvk commented Sep 24, 2026

Uh oh!

ikrommyd commented Sep 24, 2026

Uh oh!

ikrommyd commented Nov 9, 2026

Uh oh!

jorenham commented Nov 9, 2026

Uh oh!

ikrommyd commented Nov 9, 2026

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

ikrommyd commented Aug 10, 2026 •

edited

Loading

ikrommyd Aug 10, 2026 •

edited

Loading

ikrommyd Aug 11, 2026 •

edited

Loading