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_

Changed files:
tests/lax_numpy_test.py, tests/lax_numpy_reducers_test.py, jax/_src/numpy/reductions.py

diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index 23761b323..e9858fe0b 100644
--- 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))
diff --git a/tests/lax_numpy_reducers_test.py b/tests/lax_numpy_reducers_test.py
index da1a3f8ac..b8386aa41 100644
--- a/tests/lax_numpy_reducers_test.py
+++ b/tests/lax_numpy_reducers_test.py
@@ -753,6 +753,20 @@ class JaxNumpyReducerTests(jtu.JaxTestCase):
                             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
index 3f8e71bc8..2f1e6ca6d 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -5027,13 +5027,6 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     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.