Implement 2023.12 behavior for sum() and prod()

Author: asmeurer

array_api_strict/_statistical_functions.py

@@ -7,7 +7,7 @@
 )
 from ._array_object import Array
 from ._dtypes import float32, complex64
-from ._flags import requires_api_version
+from ._flags import requires_api_version, get_array_api_strict_flags
 from ._creation_functions import zeros
 from ._manipulation_functions import concat
 
@@ -89,14 +89,16 @@ def prod(
 ) -> 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 = np.float64
-        elif x.dtype == complex64:
-            dtype = np.complex128
+        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
+        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
+        # NumPy (only upcast for integral dtypes).
+        if get_array_api_strict_flags()['api_version'] < '2023.12':
+            if x.dtype == float32:
+                dtype = np.float64
+            elif x.dtype == complex64:
+                dtype = np.complex128
     else:
         dtype = dtype._np_dtype
     return Array._new(np.prod(x._array, dtype=dtype, axis=axis, keepdims=keepdims))
@@ -126,14 +128,16 @@ def sum(
 ) -> 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 = np.float64
-        elif x.dtype == complex64:
-            dtype = np.complex128
+        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
+        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
+        # NumPy (only upcast for integral dtypes).
+        if get_array_api_strict_flags()['api_version'] < '2023.12':
+            if x.dtype == float32:
+                dtype = np.float64
+            elif x.dtype == complex64:
+                dtype = np.complex128
     else:
         dtype = dtype._np_dtype
     return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims))