def test_vmaps_of_while_loop(self):

    def f(x, n):
        return lax.fori_loop(0, n, (lambda _, x: (x + 1)), x)
    (x, n) = (jnp.arange(3), jnp.arange(4))
    api.vmap(api.vmap(f, (None, 0)), (0, None))(x, n)
def test_vmaps_of_while_loop(self):

    def f(x, n):
        return lax.fori_loop(0, n, (lambda _, x: (x + 1)), x)
    (x, n) = (jnp.arange(3), jnp.arange(4))
    api.vmap(api.vmap(f, (None, 0)), (0, None))(x, n)