def testWhileWithTuple(self):
    limit = 10

    def loop_cond(state):
        (pos, _) = state
        return lax.lt(pos, limit)

    def loop_body(state):
        (pos, count) = state
        return (lax.add(pos, 1), lax.add(count, 1))

    def loop(init):
        result = lax.while_loop(loop_cond, loop_body, (init, 0))
        (_, count) = result
        return count
    cloop = api.jit(loop)
    self.assertEqual(loop(2), (limit - 2))
    self.assertEqual(cloop(2), (limit - 2))
    self.assertEqual(cloop(2), (limit - 2))
    self.assertEqual(cloop(3), (limit - 3))
def testWhileWithTuple(self):
    limit = 10

    def loop_cond(state):
        (pos, _) = state
        return lax.lt(pos, limit)

    def loop_body(state):
        (pos, count) = state
        return (lax.add(pos, 1), lax.add(count, 1))

    def loop(init):
        result = lax.while_loop(loop_cond, loop_body, (init, 0))
        (_, count) = result
        return count
    cloop = api.jit(loop)
    self.assertEqual(loop(2), (limit - 2))
    self.assertEqual(cloop(2), (limit - 2))
    self.assertEqual(cloop(2), (limit - 2))
    self.assertEqual(cloop(3), (limit - 3))