Skip to content

Commit d0d02c6

Browse files
committed
feat: add complex dtype support mean
1 parent 6d205d7 commit d0d02c6

File tree

1 file changed

+9
-3
lines changed

1 file changed

+9
-3
lines changed

src/array_api_stubs/_draft/statistical_functions.py

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ def mean(
113113
Parameters
114114
----------
115115
x: array
116-
input array. Should have a real-valued floating-point data type.
116+
input array. Should have a floating-point data type.
117117
axis: Optional[Union[int, Tuple[int, ...]]]
118118
axis or axes along which arithmetic means must be computed. By default, the mean must be computed over the entire array. If a tuple of integers, arithmetic means must be computed over multiple axes. Default: ``None``.
119119
keepdims: bool
@@ -125,17 +125,23 @@ def mean(
125125
if the arithmetic mean was computed over the entire array, a zero-dimensional array containing the arithmetic mean; otherwise, a non-zero-dimensional array containing the arithmetic means. The returned array must have the same data type as ``x``.
126126
127127
.. note::
128-
While this specification recommends that this function only accept input arrays having a real-valued floating-point data type, specification-compliant array libraries may choose to accept input arrays having an integer data type. While mixed data type promotion is implementation-defined, if the input array ``x`` has an integer data type, the returned array must have the default real-valued floating-point data type.
128+
While this specification recommends that this function only accept input arrays having a floating-point data type, specification-compliant array libraries may choose to accept input arrays having an integer data type. While mixed data type promotion is implementation-defined, if the input array ``x`` has an integer data type, the returned array must have the default real-valued floating-point data type.
129129
130130
Notes
131131
-----
132132
133133
**Special Cases**
134134
135-
Let ``N`` equal the number of elements over which to compute the arithmetic mean.
135+
Let ``N`` equal the number of elements over which to compute the arithmetic mean. For real-valued operands,
136136
137137
- If ``N`` is ``0``, the arithmetic mean is ``NaN``.
138138
- If ``x_i`` is ``NaN``, the arithmetic mean is ``NaN`` (i.e., ``NaN`` values propagate).
139+
140+
For complex floating-point operands, real-valued floating-point special cases must independently apply to the real and imaginary component operations involving real numbers. For example, let ``a = real(x_i)`` and ``b = imag(x_i)``, and
141+
142+
- If ``N`` is ``0``, the arithmetic mean is ``NaN + NaN j``.
143+
- If ``a`` is ``NaN``, the real component of the result is ``NaN``.
144+
- Similarly, if ``b`` is ``NaN``, the imaginary component of the result is ``NaN``.
139145
"""
140146

141147

0 commit comments

Comments
 (0)