Missing check in shape rule of reduce_window
`reduce_window` expects to build a reduction function of type `T -> T -> T`, where `T` is the type of the operand and both arguments are scalar. If `init_value` is not a scalar, the reduction function that is constructed is invalid, and we get a `RuntimeError`.

This behavior can be reproduced using:
```python
from jax import lax
from jax import numpy as jnp

params = { "operand": jnp.ones((4, 4), dtype=jnp.int32)
         , "init_value": jnp.zeros((1,), dtype=jnp.int32) 
         , "computation": lax.max
         , "window_dimensions": (2, 2)
         , "window_strides": (2, 2)
         , "padding": "VALID"
         , "base_dilation": (1, 1)
         , "window_dilation": (1, 1)
         }
lax.reduce_window(**params)
```
which results in
```bash
RuntimeError: Invalid argument: Reduction function must return a scalar or tuple of scalars but returns shape: s32[1]: 
This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 3c030e52f..828704fb2 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -4882,6 +4882,10 @@ def _reduce_window_shape_rule(operand, init_value, *, jaxpr, consts,
     msg = ("reduce_window got inconsistent dtypes for operand and init_value: "
            " got operand dtype {} and init_value dtype {}.")
     raise TypeError(msg.format(operand.dtype, init_value.dtype))
+  if init_value.shape != ():
+    msg = ("reduce_window expected init_value to be a scalar but init_value "
+           "has shape {}.")
+    raise TypeError(msg.format(init_value.shape))
   return _common_reduce_window_shape_rule(
     operand, window_dimensions, window_strides, padding, base_dilation,
     window_dilation)
diff --git a/tests/lax_test.py b/tests/lax_test.py
index db6fc9c02..31ab3c82f 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2031,6 +2031,22 @@ class LaxTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(ValueError, msg):
       lax.conv_general_dilated(lhs, rhs, **kwargs)
 
+  def test_reduce_window_scalar_init_value_shape_rule(self):
+    # https://github.com/google/jax/issues/4574
+    args = { "operand": np.ones((4, 4), dtype=np.int32)
+           , "init_value": np.zeros((1,), dtype=np.int32)
+           , "computation": lax.max
+           , "window_dimensions": (2, 2)
+           , "window_strides": (2, 2)
+           , "padding": "VALID"
+           , "base_dilation": (1, 1)
+           , "window_dilation": (1, 1)
+           }
+
+    msg = (r"reduce_window expected init_value to be a scalar but init_value "
+           r"has shape \\(1,\\).")
+    with self.assertRaisesRegex(TypeError, msg):
+      lax.reduce_window(**args)
 
 class LazyConstantTest(jtu.JaxTestCase):
   def _Check(self, make_const, expected):
['tests/lax_test.py', 'jax/lax/lax.py']