Repository: google/jax

Issue:
jnp.max/jnp.min allow non-scalar value for `initial`.
### Description

Hello!

I would like to report a sort of inverse feature request. Our code currently relies on the following logic working:

```python
def f(a, b):
  return jnp.max(a, axis=-1, initial=b, where=jnp.isfinite(a))

>> a = jnp.full((5,), fill_value=2)
>> b = jnp.array([1, 2, 3, 4, 5])
>> f(a, b)
Array([2, 2, 3, 4, 5], dtype=int32)
```

_Technically_ this shouldn't work since the `initial` argument of `jnp.max` is supposed to be a scalar. Numpy throws an exception `ValueError: Input object to FillWithScalar is not a scalar`, but it appears that JAX's implementation is more general and works as intended.

Can we continue to rely on this behavior?



### What jax/jaxlib version are you using?

0.4.2

### Which accelerator(s) are you using?

GPU

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


Discussion:
This behavior breaks with jit pjit merge! But magically works with the old jit.
Actually this works with both jit and pjit. The pjit merge breaks the case when `initial` has a trailing dimension of 1:

```python
def f(a, b):
  return jnp.max(a, axis=-1, initial=b, where=jnp.isfinite(a))


a = jnp.full((4, 5, 3), fill_value=2)
b = jnp.zeros((4, 5, 1))
jax.vmap(f)(a, b)  # Works with old jit, but fails with pjit.
```

The exception is `TypeError: max: arrays must have same number of dimensions, got (5, 1), (5,).`.
I think this is an unintended implementation detail, and the behavior is untested so I would probably not rely on it.  In fact, your original example looks like a bug to me, because the reduced array is the wrong size, so I'd lean toward checking the shape on `initial` and raising an error if it's not a scalar.
So I think we're going to move toward explicitly requiring `initial` to be a scalar (see https://github.com/google/jax/pull/14365#issuecomment-1426538150)

You'd still be able to effectively use broadcasted initial values by wrapping your call in `vmap`, and passing a mapped value to the `initial` argument.

Leaving this issue open to track that TODO.

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
@@ -131,11 +131,10 @@ def _reduction(a: ArrayLike, name: str, np_fun: Any, op: ReductionOp, init_val:
   else:
     result = lax.reduce(a, init_val, op, dims)
   if initial is not None:
-    # TODO(jakevdp) require initial to be a scalar in order to match the numpy API.
     initial_arr = lax.convert_element_type(initial, _asarray(a).dtype)
-    if lax.broadcast_shapes(initial_arr.shape, result.shape) != result.shape:
-      raise ValueError(f"initial value has invalid shape {initial_arr.shape} "
-                       f"for reduction with output shape {result.shape}")
+    if initial_arr.shape != ():
+      raise ValueError("initial value must be a scalar. "
+                       f"Got array of shape {initial_arr.shape}")
     result = op(initial_arr, result)
   if keepdims:
     result = lax.expand_dims(result, pos_dims)


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
@@ -287,8 +287,7 @@ def testReducerBadInitial(self, rec):
     jnp_op = getattr(jnp, rec.name)
     arr = jnp.ones((2, 3, 4))
     initial = jnp.zeros((1, 2, 3))
-    msg = (r"initial value has invalid shape \(1, 2, 3\) "
-           r"for reduction with output shape \(2, 3\)")
+    msg = r"initial value must be a scalar. Got array of shape \(1, 2, 3\)"
     with self.assertRaisesRegex(ValueError, msg):
       jnp_op(arr, axis=-1, initial=initial)