def test_disable_jit_while_loop_with_vmap(self):

    def trivial_while(y):
        return lax.while_loop((lambda x: (x < 10.0)), (lambda x: (x + 1.0)), y)
    with api.disable_jit():
        api.vmap(trivial_while)(jnp.array([3.0, 4.0]))
def test_disable_jit_while_loop_with_vmap(self):

    def trivial_while(y):
        return lax.while_loop((lambda x: (x < 10.0)), (lambda x: (x + 1.0)), y)
    with api.disable_jit():
        api.vmap(trivial_while)(jnp.array([3.0, 4.0]))