@parameterized.named_parameters(itertools.chain.from_iterable((jtu.cases_from_list(({'testcase_name': '{}_inshape={}_axis={}_dtype={}_keepdims={}'.format(rec.test_name.capitalize(), jtu.format_shape_dtype_string(shape, dtype), axis, ('None' if (out_dtype is None) else onp.dtype(out_dtype).name), keepdims), 'rng_factory': rec.rng_factory, 'shape': shape, 'dtype': dtype, 'out_dtype': out_dtype, 'onp_op': getattr(onp, rec.name), 'jnp_op': getattr(jnp, rec.name), 'axis': axis, 'keepdims': keepdims, 'inexact': rec.inexact} for shape in rec.shapes for dtype in rec.dtypes for out_dtype in ([None] + rec.dtypes) for axis in (list(range((- len(shape)), len(shape))) + [None]) for keepdims in [False, True])) for rec in JAX_REDUCER_RECORDS)))
def testReducer(self, onp_op, jnp_op, rng_factory, shape, dtype, out_dtype, axis, keepdims, inexact):
    rng = rng_factory()

    @jtu.ignore_warning(category=onp.ComplexWarning)
    @jtu.ignore_warning(category=RuntimeWarning, message='mean of empty slice.*')
    def onp_fun(x):
        x_cast = (x if (dtype != jnp.bfloat16) else x.astype(onp.float32))
        t = (out_dtype if (out_dtype != jnp.bfloat16) else onp.float32)
        return onp_op(x_cast, axis, dtype=t, keepdims=keepdims)
    onp_fun = _promote_like_jnp(onp_fun, inexact)
    jnp_fun = (lambda x: jnp_op(x, axis, dtype=out_dtype, keepdims=keepdims))
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = (lambda : [rng(shape, dtype)])
    tol_spec = {onp.float16: 0.01, onp.float32: 0.001, onp.complex64: 0.001, onp.float64: 1e-05, onp.complex128: 1e-05}
    tol = jtu.tolerance(dtype, tol_spec)
    tol = (max(tol, jtu.tolerance(out_dtype, tol_spec)) if out_dtype else tol)
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=(jnp.bfloat16 not in (dtype, out_dtype)), tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, atol=tol, rtol=tol)
@parameterized.named_parameters(itertools.chain.from_iterable((jtu.cases_from_list(({'testcase_name': '{}_inshape={}_axis={}_dtype={}_keepdims={}'.format(rec.test_name.capitalize(), jtu.format_shape_dtype_string(shape, dtype), axis, ('None' if (out_dtype is None) else onp.dtype(out_dtype).name), keepdims), 'rng_factory': rec.rng_factory, 'shape': shape, 'dtype': dtype, 'out_dtype': out_dtype, 'onp_op': getattr(onp, rec.name), 'jnp_op': getattr(jnp, rec.name), 'axis': axis, 'keepdims': keepdims, 'inexact': rec.inexact} for shape in rec.shapes for dtype in rec.dtypes for out_dtype in ([None] + rec.dtypes) for axis in (list(range((- len(shape)), len(shape))) + [None]) for keepdims in [False, True])) for rec in JAX_REDUCER_RECORDS)))
def testReducer(self, onp_op, jnp_op, rng_factory, shape, dtype, out_dtype, axis, keepdims, inexact):
    rng = rng_factory()

    @jtu.ignore_warning(category=onp.ComplexWarning)
    @jtu.ignore_warning(category=RuntimeWarning, message='mean of empty slice.*')
    def onp_fun(x):
        x_cast = (x if (dtype != jnp.bfloat16) else x.astype(onp.float32))
        t = (out_dtype if (out_dtype != jnp.bfloat16) else onp.float32)
        return onp_op(x_cast, axis, dtype=t, keepdims=keepdims)
    onp_fun = _promote_like_jnp(onp_fun, inexact)
    jnp_fun = (lambda x: jnp_op(x, axis, dtype=out_dtype, keepdims=keepdims))
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = (lambda : [rng(shape, dtype)])
    tol_spec = {onp.float16: 0.01, onp.float32: 0.001, onp.complex64: 0.001, onp.float64: 1e-05, onp.complex128: 1e-05}
    tol = jtu.tolerance(dtype, tol_spec)
    tol = (max(tol, jtu.tolerance(out_dtype, tol_spec)) if out_dtype else tol)
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=(jnp.bfloat16 not in (dtype, out_dtype)), tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True, atol=tol, rtol=tol)