def testWhileJVPViaForiLoop(self):
    f = (lambda x: lax.fori_loop(0, 3, (lambda i, x: (x * 2)), x))
    self.assertAllClose(f(2.0), 16.0, check_dtypes=False)
    self.assertAllClose(api.jvp(f, (2.0,), (1.0,)), (16.0, 8.0), check_dtypes=False)
    jtu.check_grads(f, (2.0,), order=2, modes=['fwd'])
    f = (lambda x: lax.fori_loop(0, 3, (lambda i, x: (x * (i + 1))), x))
    self.assertAllClose(f(2.0), 12.0, check_dtypes=False)
    self.assertAllClose(api.jvp(f, (2.0,), (1.0,)), (12.0, 6.0), check_dtypes=False)
    jtu.check_grads(f, (2.0,), order=2, modes=['fwd'])
def testWhileJVPViaForiLoop(self):
    f = (lambda x: lax.fori_loop(0, 3, (lambda i, x: (x * 2)), x))
    self.assertAllClose(f(2.0), 16.0, check_dtypes=False)
    self.assertAllClose(api.jvp(f, (2.0,), (1.0,)), (16.0, 8.0), check_dtypes=False)
    jtu.check_grads(f, (2.0,), order=2, modes=['fwd'])
    f = (lambda x: lax.fori_loop(0, 3, (lambda i, x: (x * (i + 1))), x))
    self.assertAllClose(f(2.0), 12.0, check_dtypes=False)
    self.assertAllClose(api.jvp(f, (2.0,), (1.0,)), (12.0, 6.0), check_dtypes=False)
    jtu.check_grads(f, (2.0,), order=2, modes=['fwd'])