def testIssue649(self):
    from jax import lax

    def body(x):
        (a, b) = x
        return (7, (b + 1))

    def cond(x):
        (a, b) = x
        return (b < 10)
    out = lax.while_loop(cond, body, (33, 4))
    self.assertEqual(out, (7, 10))
def testIssue649(self):
    from jax import lax

    def body(x):
        (a, b) = x
        return (7, (b + 1))

    def cond(x):
        (a, b) = x
        return (b < 10)
    out = lax.while_loop(cond, body, (33, 4))
    self.assertEqual(out, (7, 10))