Refactored common upcast for integral-type accumulators

Author: Micky774

jax/_src/dtypes.py

@@ -662,7 +662,7 @@ def dtype(x: Any, *, canonicalize: bool = False) -> DType:
   elif type(x) in python_scalar_dtypes:
     dt = python_scalar_dtypes[type(x)]
   elif is_type and _issubclass(x, np.generic):
-    return np.dtype(x)
+    dt = np.dtype(x)
   elif issubdtype(getattr(x, 'dtype', None), extended):
     dt = x.dtype
   else:

jax/_src/numpy/lax_numpy.py

@@ -3363,14 +3363,8 @@ def trace(a: ArrayLike, offset: int = 0, axis1: int = 0, axis2: int = 1,
     raise NotImplementedError("The 'out' argument to jnp.trace is not supported.")
   dtypes.check_user_dtype_supported(dtype, "trace")
 
+  a = asarray(a)
   a_shape = shape(a)
-  if dtype is None:
-    dtype = _dtype(a)
-    if issubdtype(dtype, integer):
-      default_int = dtypes.canonicalize_dtype(int)
-      if iinfo(dtype).bits < iinfo(default_int).bits:
-        dtype = default_int
-
   a = moveaxis(a, (axis1, axis2), (-2, -1))
 
   # Mask out the diagonal and reduce.

jax/_src/numpy/reductions.py

@@ -730,11 +730,15 @@ def cumulative_sum(
 
   axis = _canonicalize_axis(axis, x.ndim)
   dtypes.check_user_dtype_supported(dtype)
-  kind = x.dtype.kind
-  if (dtype is None and kind in {'i', 'u'}
-      and x.dtype.itemsize*8 < int(config.default_dtype_bits.value)):
-    dtype = dtypes.canonicalize_dtype(dtypes._default_types[kind])
-  x = x.astype(dtype=dtype or x.dtype)
+  if dtype is None:
+    kind = x.dtype.kind
+    if kind in {'i', 'u'}:
+      default_dtype = dtypes.canonicalize_dtype(dtypes._default_types[kind])
+      if x.dtype.itemsize*8 < int(config.default_dtype_bits.value):
+        dtype = default_dtype
+
+  if dtype:
+    x = x.astype(dtype=dtype)
   out = cumsum(x, axis=axis)
   if include_initial:
     zeros_shape = list(x.shape)

jax/experimental/array_api/_data_type_functions.py

@@ -212,18 +212,3 @@ def result_type(*arrays_and_dtypes):
   if len(dtypes) == 1:
     return dtypes[0]
   return functools.reduce(_promote_types, dtypes)
-
-
-def _promote_to_default_dtype(x):
-  if x.dtype.kind == 'b':
-    return x
-  elif x.dtype.kind == 'i':
-    return x.astype(jnp.int_)
-  elif x.dtype.kind == 'u':
-    return x.astype(jnp.uint)
-  elif x.dtype.kind == 'f':
-    return x.astype(jnp.float_)
-  elif x.dtype.kind == 'c':
-    return x.astype(jnp.complex_)
-  else:
-    raise ValueError(f"Unrecognized {x.dtype=}")

jax/experimental/array_api/_linear_algebra_functions.py

@@ -13,9 +13,6 @@
 # limitations under the License.
 
 import jax
-from jax.experimental.array_api._data_type_functions import (
-    _promote_to_default_dtype,
-)
 
 def cholesky(x, /, *, upper=False):
   """
@@ -140,7 +137,6 @@ def trace(x, /, *, offset=0, dtype=None):
   """
   Returns the sum along the specified diagonals of a matrix (or a stack of matrices) x.
   """
-  x = _promote_to_default_dtype(x)
   return jax.numpy.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1)
 
 def vecdot(x1, x2, /, *, axis=-1):

jax/experimental/array_api/_statistical_functions.py

@@ -13,9 +13,6 @@
 # limitations under the License.
 
 import jax
-from jax.experimental.array_api._data_type_functions import (
-    _promote_to_default_dtype,
-)
 
 
 def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False):
@@ -39,7 +36,6 @@ def min(x, /, *, axis=None, keepdims=False):
 
 def prod(x, /, *, axis=None, dtype=None, keepdims=False):
   """Calculates the product of input array x elements."""
-  x = _promote_to_default_dtype(x)
   return jax.numpy.prod(x, axis=axis, dtype=dtype, keepdims=keepdims)
 
 
@@ -50,7 +46,6 @@ def std(x, /, *, axis=None, correction=0.0, keepdims=False):
 
 def sum(x, /, *, axis=None, dtype=None, keepdims=False):
   """Calculates the sum of the input array x."""
-  x = _promote_to_default_dtype(x)
   return jax.numpy.sum(x, axis=axis, dtype=dtype, keepdims=keepdims)
 
 

jax/experimental/array_api/skips.txt

@@ -11,3 +11,9 @@ array_api_tests/test_special_cases.py::test_unary
 
 # fft test suite is buggy as of 83f0bcdc
 array_api_tests/test_fft.py
+
+# Pending implementation update for proper dtype promotion behavior,
+# see https://github.com/data-apis/array-api-tests/issues/234
+array_api_tests/test_statistical_functions.py::test_sum
+array_api_tests/test_statistical_functions.py::test_prod
+array_api_tests/test_linalg.py::test_trace