Restrict input dtypes in the array API statistical functions
from __future__ import annotations
+from ._dtypes import (
+ _floating_dtypes,
+ _numeric_dtypes,
+)
from ._array_object import Array
from typing import Optional, Tuple, Union
@@ -14,6 +18,8 @@ def max(
axis: Optional[Union[int, Tuple[int, ...]]] = None,
keepdims: bool = False,
) -> Array:
+ if x.dtype not in _numeric_dtypes:
+ raise TypeError("Only numeric dtypes are allowed in max")
return Array._new(np.max(x._array, axis=axis, keepdims=keepdims))
@@ -24,6 +30,8 @@ def mean(
axis: Optional[Union[int, Tuple[int, ...]]] = None,
keepdims: bool = False,
) -> Array:
+ if x.dtype not in _floating_dtypes:
+ raise TypeError("Only floating-point dtypes are allowed in mean")
return Array._new(np.mean(x._array, axis=axis, keepdims=keepdims))
@@ -34,6 +42,8 @@ def min(
axis: Optional[Union[int, Tuple[int, ...]]] = None,
keepdims: bool = False,
) -> Array:
+ if x.dtype not in _numeric_dtypes:
+ raise TypeError("Only numeric dtypes are allowed in min")
return Array._new(np.min(x._array, axis=axis, keepdims=keepdims))
@@ -44,6 +54,8 @@ def prod(
axis: Optional[Union[int, Tuple[int, ...]]] = None,
keepdims: bool = False,
) -> Array:
+ if x.dtype not in _numeric_dtypes:
+ raise TypeError("Only numeric dtypes are allowed in prod")
return Array._new(np.prod(x._array, axis=axis, keepdims=keepdims))
@@ -56,6 +68,8 @@ def std(
keepdims: bool = False,
) -> Array:
# Note: the keyword argument correction is different here
+ if x.dtype not in _floating_dtypes:
+ raise TypeError("Only floating-point dtypes are allowed in std")
return Array._new(np.std(x._array, axis=axis, ddof=correction, keepdims=keepdims))
@@ -66,6 +80,8 @@ def sum(
axis: Optional[Union[int, Tuple[int, ...]]] = None,
keepdims: bool = False,
) -> Array:
+ if x.dtype not in _numeric_dtypes:
+ raise TypeError("Only numeric dtypes are allowed in sum")
return Array._new(np.sum(x._array, axis=axis, keepdims=keepdims))
@@ -78,4 +94,6 @@ def var(
keepdims: bool = False,
) -> Array:
# Note: the keyword argument correction is different here
+ if x.dtype not in _floating_dtypes:
+ raise TypeError("Only floating-point dtypes are allowed in var")
return Array._new(np.var(x._array, axis=axis, ddof=correction, keepdims=keepdims))
