def f(rng, x):
    cond = random.bernoulli(rng)
    return lax.cond(cond, x, (lambda x: x), (jnp.abs(x) - 1.0), (lambda x: x))
def f(rng, x):
    cond = random.bernoulli(rng)
    return lax.cond(cond, x, (lambda x: x), (jnp.abs(x) - 1.0), (lambda x: x))