def testNestedWhile(self):

    def outer_loop(num):

        def cond_fun(state):
            (num, i, _) = state
            return lax.lt(i, num)

        def body_fun(state):
            (num, i, count) = state
            return (num, lax.add(i, 1), inner_loop(i, count))
        init_val = (num, 0, 0)
        (_, i, count) = lax.while_loop(cond_fun, body_fun, init_val)
        return (i, count)

    def inner_loop(i, count):

        def cond_fun(state):
            (i, j, _) = state
            return lax.le(j, i)

        def body_fun(state):
            (i, j, count) = state
            return (i, lax.add(j, 1), lax.add(count, 1))
        init_val = (i, 0, count)
        (_, _, count) = lax.while_loop(cond_fun, body_fun, init_val)
        return count
    cloop = api.jit(outer_loop)
    self.assertEqual(outer_loop(3), (3, 6))
    self.assertEqual(cloop(3), (3, 6))
    self.assertEqual(cloop(3), (3, 6))
    self.assertEqual(cloop(2), (2, 3))
    self.assertEqual(cloop(4), (4, 10))
def testNestedWhile(self):

    def outer_loop(num):

        def cond_fun(state):
            (num, i, _) = state
            return lax.lt(i, num)

        def body_fun(state):
            (num, i, count) = state
            return (num, lax.add(i, 1), inner_loop(i, count))
        init_val = (num, 0, 0)
        (_, i, count) = lax.while_loop(cond_fun, body_fun, init_val)
        return (i, count)

    def inner_loop(i, count):

        def cond_fun(state):
            (i, j, _) = state
            return lax.le(j, i)

        def body_fun(state):
            (i, j, count) = state
            return (i, lax.add(j, 1), lax.add(count, 1))
        init_val = (i, 0, count)
        (_, _, count) = lax.while_loop(cond_fun, body_fun, init_val)
        return count
    cloop = api.jit(outer_loop)
    self.assertEqual(outer_loop(3), (3, 6))
    self.assertEqual(cloop(3), (3, 6))
    self.assertEqual(cloop(3), (3, 6))
    self.assertEqual(cloop(2), (2, 3))
    self.assertEqual(cloop(4), (4, 10))