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

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

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

Remove `print`

Update tests

Remove `print`

Update tests

jax/_src/numpy/reductions.py

@@ -858,7 +858,6 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     result = lax.mul(lax.add(low_value, high_value), _lax_const(low_value, 0.5))
   else:
     raise ValueError(f"interpolation={interpolation!r} not recognized")
-  print("keepdims", keepdims, "keepdim", keepdim)
   if keepdims and keepdim:
     if q_ndim > 0:
       keepdim = [np.shape(q)[0], *keepdim]

tests/lax_numpy_reducers_test.py

@@ -662,6 +662,7 @@ def testCov(self, shape, dtype, y_shape, y_dtype, rowvar, ddof, bias, fweights,
     [dict(a_shape=a_shape, axis=axis)
       for a_shape, axis in (
         ((7,), None),
+        ((6, 7,), None),
         ((47, 7), 0),
         ((47, 7), ()),
         ((4, 101), 1),
@@ -705,12 +706,6 @@ def np_fun(*args):
                             tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker, rtol=tol)
 
-  def testQuantileKeepdims(self):
-    x = np.ones([2, 3, 4]).astype(np.float32)
-    expected = np.quantile(x, 0.5, axis=None, keepdims=True)
-    actual = jnp.quantile(x, 0.5, axis=None, keepdims=True)
-    self.assertAllClose(expected, actual, atol=0)
-
   @unittest.skipIf(not config.enable_x64.value, "test requires X64")
   @jtu.run_on_devices("cpu")  # test is for CPU float64 precision
   def testPercentilePrecision(self):
@@ -722,6 +717,7 @@ def testPercentilePrecision(self):
     [dict(a_shape=a_shape, axis=axis)
       for a_shape, axis in (
         ((7,), None),
+        ((6, 7,), None),
         ((47, 7), 0),
         ((4, 101), 1),
       )
@@ -750,12 +746,6 @@ def np_fun(*args):
                             tol=tol)
     self._CompileAndCheck(jnp_fun, args_maker, rtol=tol)
 
-  def testMedianKeepdims(self):
-    x = np.ones([2, 3, 4], dtype=np.float32)
-    expected = np.median(x, axis=None, keepdims=True)
-    actual = jnp.median(x, axis=None, keepdims=True)
-    self.assertAllClose(expected, actual, atol=0)
-
   def testMeanLargeArray(self):
     # https://github.com/google/jax/issues/15068
     raise unittest.SkipTest("test is slow, but it passes!")