def testForiLoopClosure(self):

    def count(num):

        def body_fun(i, tot):
            return lax.add(num, lax.add(tot, i))
        return lax.fori_loop(0, num, body_fun, 0)
    cfun = api.jit(count)
    self.assertEqual(count(2), (1 + (2 ** 2)))
    self.assertEqual(count(2), cfun(2))
    self.assertEqual(count(3), (3 + (3 ** 2)))
    self.assertEqual(count(3), cfun(3))
    self.assertEqual(count(4), (6 + (4 ** 2)))
    self.assertEqual(count(4), cfun(4))
def testForiLoopClosure(self):

    def count(num):

        def body_fun(i, tot):
            return lax.add(num, lax.add(tot, i))
        return lax.fori_loop(0, num, body_fun, 0)
    cfun = api.jit(count)
    self.assertEqual(count(2), (1 + (2 ** 2)))
    self.assertEqual(count(2), cfun(2))
    self.assertEqual(count(3), (3 + (3 ** 2)))
    self.assertEqual(count(3), cfun(3))
    self.assertEqual(count(4), (6 + (4 ** 2)))
    self.assertEqual(count(4), cfun(4))