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)
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)