@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_a_shape={}_axis={}_keepdims={}'.format(jtu.format_shape_dtype_string(a_shape, a_dtype), axis, keepdims), 'a_rng': jtu.rand_default(), 'a_shape': a_shape, 'a_dtype': a_dtype, 'axis': axis, 'keepdims': keepdims} for a_dtype in float_dtypes for (a_shape, axis) in (((7,), None), ((47, 7), 0), ((4, 101), 1)) for keepdims in [False, True])))
def testMedian(self, a_rng, a_shape, a_dtype, axis, keepdims):
    args_maker = (lambda : [a_rng(a_shape, a_dtype)])

    def onp_fun(*args):
        args = [(x if (jnp.result_type(x) != jnp.bfloat16) else onp.asarray(x, onp.float32)) for x in args]
        return onp.median(*args, axis=axis, keepdims=keepdims)
    jnp_fun = partial(jnp.median, axis=axis, keepdims=keepdims)
    tol_spec = {onp.float32: 0.0002, onp.float64: 5e-06}
    tol = jtu.tolerance(a_dtype, tol_spec)
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, rtol=tol)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_a_shape={}_axis={}_keepdims={}'.format(jtu.format_shape_dtype_string(a_shape, a_dtype), axis, keepdims), 'a_rng': jtu.rand_default(), 'a_shape': a_shape, 'a_dtype': a_dtype, 'axis': axis, 'keepdims': keepdims} for a_dtype in float_dtypes for (a_shape, axis) in (((7,), None), ((47, 7), 0), ((4, 101), 1)) for keepdims in [False, True])))
def testMedian(self, a_rng, a_shape, a_dtype, axis, keepdims):
    args_maker = (lambda : [a_rng(a_shape, a_dtype)])

    def onp_fun(*args):
        args = [(x if (jnp.result_type(x) != jnp.bfloat16) else onp.asarray(x, onp.float32)) for x in args]
        return onp.median(*args, axis=axis, keepdims=keepdims)
    jnp_fun = partial(jnp.median, axis=axis, keepdims=keepdims)
    tol_spec = {onp.float32: 0.0002, onp.float64: 5e-06}
    tol = jtu.tolerance(a_dtype, tol_spec)
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=False, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, rtol=tol)