Refactored common upcast for integral-type accumulators

Author: Micky774

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

@@ -26,7 +26,7 @@
 
 from jax import lax
 from jax._src import api
-from jax._src import core, config
+from jax._src import core
 from jax._src import dtypes
 from jax._src.numpy import ufuncs
 from jax._src.numpy.util import (
@@ -65,6 +65,20 @@ def _upcast_f16(dtype: DTypeLike) -> DType:
     return np.dtype('float32')
   return np.dtype(dtype)
 
+def _promote_integer_dtype(dtype: DTypeLike) -> DTypeLike:
+  # Note: NumPy always promotes to 64-bit; jax instead promotes to the
+  # default dtype as defined by dtypes.int_ or dtypes.uint.
+  if dtypes.issubdtype(dtype, np.bool_):
+    return dtypes.int_
+  elif dtypes.issubdtype(dtype, np.unsignedinteger):
+    if np.iinfo(dtype).bits < np.iinfo(dtypes.uint).bits:
+      return dtypes.uint
+  elif dtypes.issubdtype(dtype, np.integer):
+    if np.iinfo(dtype).bits < np.iinfo(dtypes.int_).bits:
+      return dtypes.int_
+  return dtype
+
+
 ReductionOp = Callable[[Any, Any], Any]
 
 def _reduction(a: ArrayLike, name: str, np_fun: Any, op: ReductionOp, init_val: ArrayLike,
@@ -103,16 +117,7 @@ def _reduction(a: ArrayLike, name: str, np_fun: Any, op: ReductionOp, init_val:
   result_dtype = dtype or dtypes.dtype(a)
 
   if dtype is None and promote_integers:
-    # Note: NumPy always promotes to 64-bit; jax instead promotes to the
-    # default dtype as defined by dtypes.int_ or dtypes.uint.
-    if dtypes.issubdtype(result_dtype, np.bool_):
-      result_dtype = dtypes.int_
-    elif dtypes.issubdtype(result_dtype, np.unsignedinteger):
-      if np.iinfo(result_dtype).bits < np.iinfo(dtypes.uint).bits:
-        result_dtype = dtypes.uint
-    elif dtypes.issubdtype(result_dtype, np.integer):
-      if np.iinfo(result_dtype).bits < np.iinfo(dtypes.int_).bits:
-        result_dtype = dtypes.int_
+    result_dtype = _promote_integer_dtype(result_dtype)
 
   result_dtype = dtypes.canonicalize_dtype(result_dtype)
 
@@ -653,7 +658,8 @@ def nanstd(a: ArrayLike, axis: Axis = None, dtype: DTypeLike | None = None, out:
 
 class CumulativeReduction(Protocol):
   def __call__(self, a: ArrayLike, axis: Axis = None,
-               dtype: DTypeLike | None = None, out: None = None) -> Array: ...
+               dtype: DTypeLike | None = None, out: None = None,
+               promote_integers: bool = False) -> Array: ...
 
 
 # TODO(jakevdp): should we change these semantics to match those of numpy?
@@ -667,12 +673,17 @@ def _make_cumulative_reduction(np_reduction: Any, reduction: Callable[..., Array
   @implements(np_reduction, skip_params=['out'],
           lax_description=CUML_REDUCTION_LAX_DESCRIPTION)
   def cumulative_reduction(a: ArrayLike, axis: Axis = None,
-                           dtype: DTypeLike | None = None, out: None = None) -> Array:
-    return _cumulative_reduction(a, _ensure_optional_axes(axis), dtype, out)
+                           dtype: DTypeLike | None = None, out: None = None,
+                           promote_integers: bool = False) -> Array:
+    return _cumulative_reduction(
+      a, _ensure_optional_axes(axis), dtype,
+      out, promote_integers=promote_integers
+    )
 
-  @partial(api.jit, static_argnames=('axis', 'dtype'))
+  @partial(api.jit, static_argnames=('axis', 'dtype', 'promote_integers'))
   def _cumulative_reduction(a: ArrayLike, axis: Axis = None,
-                            dtype: DTypeLike | None = None, out: None = None) -> Array:
+                            dtype: DTypeLike | None = None, out: None = None,
+                            promote_integers: bool = False) -> Array:
     check_arraylike(np_reduction.__name__, a)
     if out is not None:
       raise NotImplementedError(f"The 'out' argument to jnp.{np_reduction.__name__} "
@@ -691,11 +702,15 @@ def _cumulative_reduction(a: ArrayLike, axis: Axis = None,
     if fill_nan:
       a = _where(lax_internal._isnan(a), _lax_const(a, fill_value), a)
 
-    if not dtype and dtypes.dtype(a) == np.bool_:
+    result_type = dtype or dtypes.dtype(a)
+    if result_type == np.bool_:
       dtype = dtypes.canonicalize_dtype(dtypes.int_)
+    elif dtype is None and promote_integers:
+      dtype = _promote_integer_dtype(result_type)
     if dtype:
       a = lax.convert_element_type(a, dtype)
 
+
     return reduction(a, axis)
 
   return cumulative_reduction
@@ -730,12 +745,7 @@ 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)
-  out = cumsum(x, axis=axis)
+  out = cumsum(x, axis=axis, dtype=dtype, promote_integers=True)
   if include_initial:
     zeros_shape = list(x.shape)
     zeros_shape[axis] = 1

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