MAINT: simplify cumulative_prod

Author: ev-br

array_api_strict/_dtypes.py

@@ -127,34 +127,6 @@ def __hash__(self):
 }
 
 
-def _bit_width(dtype):
-    """The bit width of an integer dtype"""
-    if dtype == int8 or dtype == uint8:
-        return 8
-    elif dtype == int16 or dtype == uint16:
-        return 16
-    elif dtype == int32 or dtype == uint32:
-        return 32
-    elif dtype == int64 or dtype == uint64:
-        return 64
-    else:
-        raise ValueError(f"_bit_width: {dtype = } not understood.")
-
-
-def _get_unsigned_from_signed(dtype):
-    """Return an unsigned integral dtype to match the input dtype."""
-    if dtype == int8:
-        return uint8
-    elif dtype == int16:
-        return uint16
-    elif dtype == int32:
-        return uint32
-    elif dtype == int64:
-        return uint64
-    else:
-        raise ValueError(f"_unsigned_from_signed: {dtype = } not understood.")
-
-
 # Note: the spec defines a restricted type promotion table compared to NumPy.
 # In particular, cross-kind promotions like integer + float or boolean +
 # integer are not allowed, even for functions that accept both kinds.

array_api_strict/_statistical_functions.py

@@ -5,14 +5,11 @@
     _real_numeric_dtypes,
     _floating_dtypes,
     _numeric_dtypes,
-    _integer_dtypes
 )
-from . import _dtypes
-from . import _info
 from ._array_object import Array
 from ._dtypes import float32, complex64
 from ._flags import requires_api_version, get_array_api_strict_flags
-from ._creation_functions import zeros
+from ._creation_functions import zeros, ones
 from ._manipulation_functions import concat
 
 from typing import TYPE_CHECKING
@@ -65,23 +62,8 @@ def cumulative_prod(
     if x.ndim == 0:
         raise ValueError("Only ndim >= 1 arrays are allowed in cumulative_prod")
 
-    # TODO: either all this is done by numpy's cumprod (?), or cumulative_sum should follow the same dance.
-    if dtype is None:
-        if x.dtype in _integer_dtypes:
-            default_int = _info.__array_namespace_info__().default_dtypes()["integral"]
-            if _dtypes._bit_width(x.dtype) < _dtypes._bit_width(default_int):
-                if x.dtype in _dtypes._unsigned_integer_dtypes:
-                    # find the unsigned integer of the same width as `default_int`
-                    dtype = _dtypes._get_unsigned_from_signed(default_int)
-                else:
-                    dtype = default_int
-            else:
-                dtype = x.dtype
-        else:
-            dtype = x.dtype
-    else:
-        if x.dtype != dtype:
-            x = xp.astype(dtype)
+    if dtype is not None:
+        dtype = dtype._np_dtype
 
     if axis is None:
         if x.ndim > 1:
@@ -92,8 +74,8 @@ def cumulative_prod(
     if include_initial:
         if axis < 0:
             axis += x.ndim
-        x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=dtype), x], axis=axis)
-    return Array._new(np.cumprod(x._array, axis=axis, dtype=dtype._np_dtype), device=x.device)
+        x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
+    return Array._new(np.cumprod(x._array, axis=axis, dtype=dtype), device=x.device)
 
 
 def max(