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.

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
@@ -72,7 +72,7 @@ def _reduction(a: ArrayLike, name: str, np_fun: Any, op: ReductionOp, init_val:
                axis: Axis = None, dtype: DTypeLike = None, out: None = None,
                keepdims: bool = False, initial: Optional[ArrayLike] = None,
                where_: Optional[ArrayLike] = None,
-               parallel_reduce: Optional[Callable[..., ArrayLike]] = None,
+               parallel_reduce: Optional[Callable[..., Array]] = None,
                promote_integers: bool = False) -> Array:
   bool_op = bool_op or op
   # Note: we must accept out=None as an argument, because numpy reductions delegate to
@@ -131,7 +131,12 @@ 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:
-    result = op(lax.convert_element_type(initial, _asarray(a).dtype), result)
+    # 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}")
+    result = op(initial_arr, result)
   if keepdims:
     result = lax.expand_dims(result, pos_dims)
   return lax.convert_element_type(result, dtype or result_dtype)


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
@@ -282,6 +282,16 @@ def np_fun(x):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker, rtol=tol, atol=tol)
 
+  @jtu.sample_product(rec = JAX_REDUCER_INITIAL_RECORDS)
+  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\)")
+    with self.assertRaisesRegex(ValueError, msg):
+      jnp_op(arr, axis=-1, initial=initial)
+
   @parameterized.parameters(itertools.chain.from_iterable(
     jtu.sample_product_testcases(
       [dict(name=rec.name, rng_factory=rec.rng_factory, inexact=rec.inexact)],