Return nan from jnp.median on an empty reduction axis

[teddytennant]

jnp.median raises an internal gather error when the reduction axis is empty, where NumPy returns nan:

>>> np.median(np.empty((0, 4), np.float32), axis=0)
array([nan, nan, nan, nan], dtype=float32)
>>> jnp.median(jnp.empty((0, 4), jnp.float32), axis=0)
TypeError: Slice size at index 0 in gather op is out of range, must be within [0, 0 + 1), got 1.

Same for axis=None, tuple axes, negative axes and keepdims=True. jnp.nanmedian already returns nan, so the two disagree today.

#39475 added an empty-reduction-axis early return to _quantile, but gated it on squash_nans, which only the nan* entry points set. median goes through quantile(a, 0.5, method='midpoint') with squash_nans=False and falls through to the sort and gather, where the gather shape rule rejects a slice size of 1 against a dimension of size 0.

This adds a nan_if_empty flag to _quantile and has median pass it. median now calls _quantile directly, because the flag is not something jnp.quantile should expose. nanmedian is unchanged: it routes through nanquantile, which already has the right empty behavior. One side effect worth flagging: the error for median(x, overwrite_input=True) and median(x, out=...) now says jax.numpy.median does not support ... instead of jax.numpy.quantile does not support ....

quantile and percentile are deliberately left alone. You asked on #39475 for them to be covered too, but NumPy only returns nan for median; for quantile and percentile NumPy raises IndexError, and JAX raises as well, so the only thing left to change there is which error, which is a design call rather than a bug fix. Happy to fold that into this PR if you want it, IndexError or otherwise.

Tests: testMedianEmptyReduction in tests/lax_numpy_reducers_test.py, checked against NumPy over 1-D, 2-D and 3-D empty inputs, axis=None, int, negative and tuple axes, keepdims both ways, eager and jit, plus two cases where the array is empty but the reduction axis is not.

$ JAX_NUM_GENERATED_CASES=1000 pytest -q tests/lax_numpy_reducers_test.py -k EmptyReduction
54 passed, 46856 deselected

$ pytest -q tests/lax_numpy_reducers_test.py
710 passed, 4 skipped          # 700 passed, 4 skipped before this change

$ JAX_ENABLE_X64=1 pytest -q tests/lax_numpy_reducers_test.py
715 passed, 3 skipped          # 705 passed, 3 skipped before this change

With the source change reverted and the test kept, 20 of the 22 cases fail; the 2 that pass are the non-empty-reduction-axis controls. CPU only, I have no GPU or TPU to test on.

[jakevdp]

I wonder if we really want to mimic NumPy's behavior here? When would it be useful to silently get an array of NaNs instead of getting a hard error during tracing? The current behavior seems much safer to me.