@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': jtu.format_test_name_suffix('nan_to_num', [shape], [dtype]), 'rng_factory': jtu.rand_some_inf_and_nan, 'shape': shape, 'dtype': dtype} for shape in all_shapes for dtype in inexact_dtypes)))
def testNanToNum(self, rng_factory, shape, dtype):
    rng = rng_factory()
    dtype = onp.dtype(dtypes.canonicalize_dtype(dtype)).type

    def onp_fun(x):
        if (dtype == jnp.bfloat16):
            x = onp.where(onp.isnan(x), dtype(0), x)
            x = onp.where(onp.isposinf(x), jnp.finfo(dtype).max, x)
            x = onp.where(onp.isneginf(x), jnp.finfo(dtype).min, x)
            return x
        else:
            return onp.nan_to_num(x).astype(dtype)
    args_maker = (lambda : [rng(shape, dtype)])
    check_dtypes = (shape is not jtu.PYTHON_SCALAR_SHAPE)
    self._CheckAgainstNumpy(onp_fun, jnp.nan_to_num, args_maker, check_dtypes=check_dtypes)
    self._CompileAndCheck(jnp.nan_to_num, args_maker, check_dtypes=check_dtypes)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': jtu.format_test_name_suffix('nan_to_num', [shape], [dtype]), 'rng_factory': jtu.rand_some_inf_and_nan, 'shape': shape, 'dtype': dtype} for shape in all_shapes for dtype in inexact_dtypes)))
def testNanToNum(self, rng_factory, shape, dtype):
    rng = rng_factory()
    dtype = onp.dtype(dtypes.canonicalize_dtype(dtype)).type

    def onp_fun(x):
        if (dtype == jnp.bfloat16):
            x = onp.where(onp.isnan(x), dtype(0), x)
            x = onp.where(onp.isposinf(x), jnp.finfo(dtype).max, x)
            x = onp.where(onp.isneginf(x), jnp.finfo(dtype).min, x)
            return x
        else:
            return onp.nan_to_num(x).astype(dtype)
    args_maker = (lambda : [rng(shape, dtype)])
    check_dtypes = (shape is not jtu.PYTHON_SCALAR_SHAPE)
    self._CheckAgainstNumpy(onp_fun, jnp.nan_to_num, args_maker, check_dtypes=check_dtypes)
    self._CompileAndCheck(jnp.nan_to_num, args_maker, check_dtypes=check_dtypes)