array API: add cumulative_sum and cumulative_prod#3731
Conversation
zcbenz
left a comment
zcbenz
left a comment
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.
|
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
left a comment
zcbenz
left a comment
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.
|
I agree with @zcbenz 's comments and I want to add that this implementation is kind of inefficient and in a hidden way. Can we collect a bit of information regarding the use of this particular API specifically the Just to be clear I foresee code being written like mx.cumulative_sum(x, axis=-1, include_initial=True)[..., :-1]which is a very inefficient way of writing mx.cumsum(x, axis=-1, inclusive=False) |
|
Thanks for the feedback @angeloskath. To answer your question: this is purely for compatibility with the Python array API spec, which requires You're right that the current approach (cumsum + zeros + concatenate) is three ops and can lead to the inefficient usage pattern you showed. To address that, we could implement |
3 participants
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.