def testForiLoopBatchedIssue1190(self):
    cond_fun = (lambda carry: (carry[0] < 4))
    body_fun = (lambda carry: ((carry[0] + 1), (carry[1] + 1)))
    f = (lambda x: lax.while_loop(cond_fun, body_fun, (0, x)))
    jaxpr = api.make_jaxpr(api.vmap(f))(jnp.arange(3))
    eqn = jaxpr.jaxpr.eqns[0]
    self.assertIs(eqn.primitive, lax.while_p)
    self.assertEqual(eqn.params['cond_jaxpr'].in_avals[0].shape, ())
def testForiLoopBatchedIssue1190(self):
    cond_fun = (lambda carry: (carry[0] < 4))
    body_fun = (lambda carry: ((carry[0] + 1), (carry[1] + 1)))
    f = (lambda x: lax.while_loop(cond_fun, body_fun, (0, x)))
    jaxpr = api.make_jaxpr(api.vmap(f))(jnp.arange(3))
    eqn = jaxpr.jaxpr.eqns[0]
    self.assertIs(eqn.primitive, lax.while_p)
    self.assertEqual(eqn.params['cond_jaxpr'].in_avals[0].shape, ())