improve jnp.mean / jnp.sum / ... error message for out-of-bounds axis index

[mattjj]

jnp.mean(0, axis=0)

Before:

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/james10.py", line 4, in <module>
    jnp.mean(0, axis=0)
  File "/usr/local/google/home/mattjj/packages/jax/jax/_src/numpy/reductions.py", line 844, in mean
    return _mean(a, _ensure_optional_axes(axis), dtype, out, keepdims,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/mattjj/packages/jax/jax/_src/numpy/reductions.py", line 871, in _mean
    normalizer = core.dimension_as_value(_axis_size(a, axis))
                                         ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/mattjj/packages/jax/jax/_src/numpy/reductions.py", line 782, in _axis_size
    size *= maybe_named_axis(a, lambda i: a_shape[i], lambda name: lax.psum(1, name))
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/mattjj/packages/jax/jax/_src/numpy/reductions.py", line 782, in <lambda>
    size *= maybe_named_axis(a, lambda i: a_shape[i], lambda name: lax.psum(1, name))
                                          ~~~~~~~^^^
IndexError: tuple index out of range

After:

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/james10.py", line 4, in <module>
    jnp.mean(0, axis=0)
  File "/usr/local/google/home/mattjj/packages/jax/jax/_src/numpy/reductions.py", line 861, in mean
    axis = _ensure_optional_axes(axis, core.get_aval(a))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/google/home/mattjj/packages/jax/jax/_src/numpy/reductions.py", line 216, in _ensure_optional_axes
    raise ValueError(f"'axis' argument of {x} is out-of-bounds for array of "
ValueError: 'axis' argument of 0 is out-of-bounds for array of rank 0 (type i32[])