def testIssue1263(self):

    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 body_fn(i, state):
        (rng, x) = state
        (key, subkey) = random.split(rng)
        return (key, f(subkey, x))

    def g(rng, x):
        return lax.fori_loop(0, 10, body_fn, (rng, x))
    api.vmap(g)(random.split(random.PRNGKey(0), 3), jnp.ones((3, 4)))
def testIssue1263(self):

    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 body_fn(i, state):
        (rng, x) = state
        (key, subkey) = random.split(rng)
        return (key, f(subkey, x))

    def g(rng, x):
        return lax.fori_loop(0, 10, body_fn, (rng, x))
    api.vmap(g)(random.split(random.PRNGKey(0), 3), jnp.ones((3, 4)))