def testForiLoopBasic(self):

    def body_fun(i, tot):
        return lax.add(tot, i)

    def count(num):
        return lax.fori_loop(0, num, body_fun, 0)
    self.assertEqual(count(2), 1)
    self.assertEqual(count(3), 3)
    self.assertEqual(count(4), 6)
    for args_maker in [(lambda : [2]), (lambda : [3]), (lambda : [4])]:
        self._CompileAndCheck(count, args_maker)
def testForiLoopBasic(self):

    def body_fun(i, tot):
        return lax.add(tot, i)

    def count(num):
        return lax.fori_loop(0, num, body_fun, 0)
    self.assertEqual(count(2), 1)
    self.assertEqual(count(3), 3)
    self.assertEqual(count(4), 6)
    for args_maker in [(lambda : [2]), (lambda : [3]), (lambda : [4])]:
        self._CompileAndCheck(count, args_maker)