Fix axis=None, keepdims=True for jnp.quantile and jnp.median

[james77777778]

This PR fixes the incorrect behavior:

>>> x = np.ones([2, 3, 4])
>>> np.median(x, axis=None, keepdims=True).shape
(1, 1, 1)
>>> jnp.median(x, axis=None, keepdims=True).shape
(1,)  # mismatched!
>>> np.quantile(x, 0.5, axis=None, keepdims=True).shape
(1, 1, 1)
>>> jnp.quantile(x, 0.5, axis=None, keepdims=True).shape
(1,)  # mismatched!

The corresponding tests have been included

[jakevdp]

Good catch, thanks for the fix! One minor suggestion: rather than new specific test cases, we could add relevant cases to the old tests. For example here: https://github.com/google/jax/blob/64bd95ded50fc647d6e80a55ab1cb3baef368cf6/tests/lax_numpy_reducers_test.py#L664
and here: https://github.com/google/jax/blob/64bd95ded50fc647d6e80a55ab1cb3baef368cf6/tests/lax_numpy_reducers_test.py#L718
We could add another case that looks something like ((6, 7), None) and that would lead to this case being covered for our runs with larger JAX_NUM_GENERATED_CASES. What do you think?

[james77777778]

We could add another case that looks something like ((6, 7), None) and that would lead to this case being covered for our runs with larger JAX_NUM_GENERATED_CASES. What do you think?

I'm not familar with the test configuration of JAX. Thanks for the suggestion.
The tests have been updated and the new tests should be included using JAX_NUM_GENERATED_CASES=40

[jakevdp]

Looks great, thanks!

Last thing: could you please squash the changes into a single commit? See https://jax.readthedocs.io/en/latest/contributing.html#single-change-commits-and-pull-requests. Thanks!

[james77777778]

Last thing: could you please squash the changes into a single commit? See https://jax.readthedocs.io/en/latest/contributing.html#single-change-commits-and-pull-requests. Thanks!

It should be a one-commit PR right now. Thanks for the tip.

[james77777778]

Kindly ping @jakevdp
Is there any unresolved issue? Happy to update.

[jakevdp]

Thanks!