Repository: google/jax

Issue:
jnp.mean and jnp.std crash on large arrays.
### Description


I have a large array of shape float32[16,32,1280,4096].  When I call jnp.std on this array, jit crashes and fails with the following error:  "Python int 2684354560 too large to convert to int32".

Clearly the implementation of jnp.std is dividing by the number of elements in the array, which in this case does not fit into a 32-bit integer.  

Target platform: pjit on TPU. 


### What jax/jaxlib version are you using?

v0.4.2

### Which accelerator(s) are you using?

TPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
Thanks, Delesley! I presume we should convert to float earlier, which the division will do anyway.

Resolving patch:
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -433,8 +433,9 @@ def _var(a: ArrayLike, axis: Axis = None, dtype: DTypeLike = None,
   a = lax_internal.asarray(a).astype(computation_dtype)
   a_mean = mean(a, axis, dtype=computation_dtype, keepdims=True, where=where)
   centered = lax.sub(a, a_mean)
-  if dtypes.issubdtype(centered.dtype, np.complexfloating):
+  if dtypes.issubdtype(computation_dtype, np.complexfloating):
     centered = lax.real(lax.mul(centered, lax.conj(centered)))
+    computation_dtype = centered.dtype  # avoid casting to complex below.
   else:
     centered = lax.square(centered)
 
@@ -443,13 +444,13 @@ def _var(a: ArrayLike, axis: Axis = None, dtype: DTypeLike = None,
       normalizer = core.dimension_as_value(np.size(a))
     else:
       normalizer = core.dimension_as_value(_axis_size(a, axis))
+    normalizer = lax.convert_element_type(normalizer, computation_dtype)
   else:
-    normalizer = sum(_broadcast_to(where, np.shape(a)), axis, dtype=dtype, keepdims=keepdims)
-  normalizer = normalizer - ddof
-
-  result = sum(centered, axis, keepdims=keepdims, where=where)
-  result = lax.div(result, lax.convert_element_type(normalizer, result.dtype))
-  return lax.convert_element_type(result, dtype)
+    normalizer = sum(_broadcast_to(where, np.shape(a)), axis,
+                     dtype=computation_dtype, keepdims=keepdims)
+  normalizer = lax.sub(normalizer, lax.convert_element_type(ddof, computation_dtype))
+  result = sum(centered, axis, dtype=computation_dtype, keepdims=keepdims, where=where)
+  return lax.div(result, normalizer).astype(dtype)
 
 
 def _var_promote_types(a_dtype: DTypeLike, dtype: DTypeLike) -> Tuple[DType, DType]:
@@ -486,6 +487,8 @@ def _std(a: ArrayLike, axis: Axis = None, dtype: DTypeLike = None,
          where: Optional[ArrayLike] = None) -> Array:
   check_arraylike("std", a)
   dtypes.check_user_dtype_supported(dtype, "std")
+  if dtype is not None and not dtypes.issubdtype(dtype, np.inexact):
+    raise ValueError(f"dtype argument to jnp.std must be inexact; got {dtype}")
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.std is not supported.")
   return lax.sqrt(var(a, axis=axis, dtype=dtype, ddof=ddof, keepdims=keepdims, where=where))


Test patch:
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -753,6 +753,20 @@ def np_fun(*args):
                             tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker, rtol=tol)
 
+  def testMeanLargeArray(self):
+    # https://github.com/google/jax/issues/15068
+    raise unittest.SkipTest("test is slow, but it passes!")
+    x = jnp.ones((16, 32, 1280, 4096), dtype='int8')
+    self.assertEqual(1.0, jnp.mean(x))
+    self.assertEqual(1.0, jnp.mean(x, where=True))
+
+  def testStdLargeArray(self):
+    # https://github.com/google/jax/issues/15068
+    raise unittest.SkipTest("test is slow, but it passes!")
+    x = jnp.ones((16, 32, 1280, 4096), dtype='int8')
+    self.assertEqual(0.0, jnp.std(x))
+    self.assertEqual(0.0, jnp.std(x, where=True))
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -5027,13 +5027,6 @@ def testFromString(self):
     actual = jnp.fromstring(s, sep=',', dtype=int)
     self.assertArraysEqual(expected, actual)
 
-  def testMeanLargeArray(self):
-    # https://github.com/google/jax/issues/15068
-    raise unittest.SkipTest("test is slow, but it passes!")
-    x = jnp.ones((16, 32, 1280, 4096), dtype='int8')
-    self.assertEqual(1.0, jnp.mean(x))
-    self.assertEqual(1.0, jnp.mean(x, where=True))
-
 
 # Most grad tests are at the lax level (see lax_test.py), but we add some here
 # as needed for e.g. particular compound ops of interest.