diff options
author | S. Solomon Darnell | 2025-03-28 21:52:21 -0500 |
---|---|---|
committer | S. Solomon Darnell | 2025-03-28 21:52:21 -0500 |
commit | 4a52a71956a8d46fcb7294ac71734504bb09bcc2 (patch) | |
tree | ee3dc5af3b6313e921cd920906356f5d4febc4ed /.venv/lib/python3.12/site-packages/numpy/array_api/_statistical_functions.py | |
parent | cc961e04ba734dd72309fb548a2f97d67d578813 (diff) | |
download | gn-ai-4a52a71956a8d46fcb7294ac71734504bb09bcc2.tar.gz |
Diffstat (limited to '.venv/lib/python3.12/site-packages/numpy/array_api/_statistical_functions.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/numpy/array_api/_statistical_functions.py | 122 |
1 files changed, 122 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/numpy/array_api/_statistical_functions.py b/.venv/lib/python3.12/site-packages/numpy/array_api/_statistical_functions.py new file mode 100644 index 00000000..98e31b51 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/numpy/array_api/_statistical_functions.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +from ._dtypes import ( + _real_floating_dtypes, + _real_numeric_dtypes, + _numeric_dtypes, +) +from ._array_object import Array +from ._dtypes import float32, float64, complex64, complex128 + +from typing import TYPE_CHECKING, Optional, Tuple, Union + +if TYPE_CHECKING: + from ._typing import Dtype + +import numpy as np + + +def max( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, +) -> Array: + if x.dtype not in _real_numeric_dtypes: + raise TypeError("Only real numeric dtypes are allowed in max") + return Array._new(np.max(x._array, axis=axis, keepdims=keepdims)) + + +def mean( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, +) -> Array: + if x.dtype not in _real_floating_dtypes: + raise TypeError("Only real floating-point dtypes are allowed in mean") + return Array._new(np.mean(x._array, axis=axis, keepdims=keepdims)) + + +def min( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: bool = False, +) -> Array: + if x.dtype not in _real_numeric_dtypes: + raise TypeError("Only real numeric dtypes are allowed in min") + return Array._new(np.min(x._array, axis=axis, keepdims=keepdims)) + + +def prod( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + dtype: Optional[Dtype] = None, + keepdims: bool = False, +) -> Array: + if x.dtype not in _numeric_dtypes: + raise TypeError("Only numeric dtypes are allowed in prod") + # Note: sum() and prod() always upcast for dtype=None. `np.prod` does that + # for integers, but not for float32 or complex64, so we need to + # special-case it here + if dtype is None: + if x.dtype == float32: + dtype = float64 + elif x.dtype == complex64: + dtype = complex128 + return Array._new(np.prod(x._array, dtype=dtype, axis=axis, keepdims=keepdims)) + + +def std( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, +) -> Array: + # Note: the keyword argument correction is different here + if x.dtype not in _real_floating_dtypes: + raise TypeError("Only real floating-point dtypes are allowed in std") + return Array._new(np.std(x._array, axis=axis, ddof=correction, keepdims=keepdims)) + + +def sum( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + dtype: Optional[Dtype] = None, + keepdims: bool = False, +) -> Array: + if x.dtype not in _numeric_dtypes: + raise TypeError("Only numeric dtypes are allowed in sum") + # Note: sum() and prod() always upcast for dtype=None. `np.sum` does that + # for integers, but not for float32 or complex64, so we need to + # special-case it here + if dtype is None: + if x.dtype == float32: + dtype = float64 + elif x.dtype == complex64: + dtype = complex128 + return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims)) + + +def var( + x: Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + correction: Union[int, float] = 0.0, + keepdims: bool = False, +) -> Array: + # Note: the keyword argument correction is different here + if x.dtype not in _real_floating_dtypes: + raise TypeError("Only real floating-point dtypes are allowed in var") + return Array._new(np.var(x._array, axis=axis, ddof=correction, keepdims=keepdims)) |