Skip to content

Add specification for computing the cumulative sum #653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Dec 14, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions spec/draft/API_specification/statistical_functions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ Objects in API
:toctree: generated
:template: method.rst

cumulative_sum
max
mean
min
Expand Down
60 changes: 59 additions & 1 deletion src/array_api_stubs/_draft/statistical_functions.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,67 @@
__all__ = ["max", "mean", "min", "prod", "std", "sum", "var"]
__all__ = ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]


from ._types import Optional, Tuple, Union, array, dtype


def cumulative_sum(
x: array,
/,
*,
axis: Optional[int] = None,
dtype: Optional[dtype] = None,
include_initial: bool = False,
) -> array:
"""
Calculates the cumulative sum of elements in the input array ``x``.

Parameters
----------
x: array
input array. Should have a numeric data type.
axis: Optional[int]
axis along which a cumulative sum must be computed. If ``axis`` is negative, the function must determine the axis along which to compute a cumulative sum by counting from the last dimension.

If ``x`` is a one-dimensional array, providing an ``axis`` is optional; however, if ``x`` has more than one dimension, providing an ``axis`` is required.
dtype: Optional[dtype]
data type of the returned array. If ``None``,

- if the default data type corresponding to the data type "kind" (integer, real-valued floating-point, or complex floating-point) of ``x`` has a smaller range of values than the data type of ``x`` (e.g., ``x`` has data type ``int64`` and the default data type is ``int32``, or ``x`` has data type ``uint64`` and the default data type is ``int64``), the returned array must have the same data type as ``x``.

- if the default data type corresponding to the data type "kind" of ``x`` has the same or a larger range of values than the data type of ``x``,

- if ``x`` has a real-valued floating-point data type, the returned array must have the default real-valued floating-point data type.
- if ``x`` has a complex floating-point data type, the returned array must have the default complex floating-point data type.
- if ``x`` has a signed integer data type (e.g., ``int16``), the returned array must have the default integer data type.
- if ``x`` has an unsigned integer data type (e.g., ``uint16``), the returned array must have an unsigned integer data type having the same number of bits as the default integer data type (e.g., if the default integer data type is ``int32``, the returned array must have a ``uint32`` data type).

If the data type (either specified or resolved) differs from the data type of ``x``, the input array should be cast to the specified data type before computing the sum. Default: ``None``.

.. note::
keyword argument is intended to help prevent data type overflows.

include_initial: bool
boolean indicating whether to include the initial value as the first value in the output. By convention, the initial value must be the additive identity (i.e., zero). Default: ``False``.

Returns
-------
out: array
an array containing the cumulative sums. The returned array must have a data type as described by the ``dtype`` parameter above.

Let ``N`` be the size of the axis along which to compute the cumulative sum. The returned array must have a shape determined according to the following rules:

- if ``include_initial`` is ``True``, the returned array must have the same shape as ``x``, except the size of the axis along which to compute the cumulative sum must be ``N+1``.
- if ``include_initial`` is ``False``, the returned array must have the same shape as ``x``.

Notes
-----

**Special Cases**

For both real-valued and complex floating-point operands, special cases must be handled as if the operation is implemented by successive application of :func:`~array_api.add`.
"""


def max(
x: array,
/,
Expand Down