array API: add cumulative_sum and cumulative_prod#3731
Open
katlun-lgtm wants to merge 5 commits into
Open
Conversation
4 tasks
zcbenz
requested changes
Jun 21, 2026
zcbenz
left a comment
zcbenz
left a comment
Collaborator
There was a problem hiding this comment.
It should be done by making cumsum support the args of cumulative_sum and then making cumulative_sum a pure alias.
katlun-lgtm
force-pushed
the
array-api-cumulative
branch
4 times, most recently
from
959d88d to
041ae18
Compare
These are the Array API standard equivalents of cumsum/cumprod with three key differences that justify the separate names: 1. axis=None (default) flattens the input first; cumsum/cumprod require an explicit axis. 2. include_initial=True prepends the identity element (0 for sum, 1 for prod) so the output length along axis is len+1. This matches the Array API spec's include_initial parameter and has no equivalent in cumsum/cumprod. 3. dtype parameter casts the input before accumulating, matching NumPy 2.0 / Array API behaviour. Docs and tests included. Part of the array API split from ml-explore#3684.
katlun-lgtm
force-pushed
the
array-api-cumulative
branch
from
041ae18 to
a75dd87
Compare
Without nb::sig(), nanobind's auto-generated __doc__ line 1 is parsed as RST by Sphinx, causing 'Inline emphasis start-string without end- string' warnings from the *, keyword-only separator. With nb::sig() present, Sphinx recognises line 1 as a Python function signature and skips RST markup parsing of that line.
The pure-alias approach (m.attr = m.attr) caused Sphinx to see
cumulative_sum.__doc__ starting with 'cumsum(...)' — a name mismatch
that prevented signature stripping, leaving '*,' to be parsed as RST
emphasis → 'Inline emphasis start-string without end-string'.
Separate bindings with their own nb::sig('def cumulative_sum(...)') fix
this: Sphinx sees the correct function name, strips the signature line,
and processes only the plain docstring body as RST.
cumsum and cumprod are restored to upstream-identical implementations.
cumulative_sum and cumulative_prod add dtype and include_initial params
as required by the array API standard.
Collaborator
|
Sorry if I was not clear but the correct way to fix this is to add the |
…e aliases Per zcbenz review: extend cumsum and cumprod with dtype and include_initial params, then expose cumulative_sum and cumulative_prod as pure aliases (m.attr = m.attr). Also remove cumulative_sum/cumulative_prod from ops.rst autosummary — all other pure aliases (empty, pow, matrix_transpose) are intentionally absent from ops.rst; listing them caused Sphinx to encounter a name mismatch in the shared __doc__ (__doc__ starts with 'cumsum(...)' but Sphinx is documenting 'cumulative_sum') which prevented signature stripping and triggered an RST emphasis warning on '*, '.
zcbenz
reviewed
Jun 23, 2026
zcbenz
left a comment
zcbenz
left a comment
Collaborator
There was a problem hiding this comment.
The implementation of new parameter should live in the C++ version of the ops not the python wrapper.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Focused split from #3684. Adds
cumulative_sumandcumulative_prodas the Array API standard equivalents ofcumsum/cumprod.Why separate names? (addressing @zcbenz's concern in #3684)
These are intentionally separate from
cumsum/cumprodbecause they have different semantics on three axes:cumsum/cumprodcumulative_sum/cumulative_prodaxis=Noneinclude_initial=Trueprepends identity (0 or 1), output length = input length + 1dtypeparaminclude_initialis specified by the Array API standard and enables patterns like prefix-sum where you need the neutral element at position 0. It cannot be retrofitted ontocumsum/cumprodwithout a breaking change.Files changed
python/src/ops.cpp— C++ lambda registrationsdocs/src/python/ops.rst— alphabetical entriespython/tests/test_ops.py—test_cumulative_sum_prodRelated
Split from #3684 per reviewer feedback from @zcbenz.