def testWhileJVPWithGrowingNonzeroTangents(self):
    rng = np.random.RandomState(0)

    def cond(state):
        (i, x, y, z) = state
        return (i < 2)

    def body(state):
        (i, x, y, z) = state
        y = (x * x)
        z = (y * y)
        return ((i + 1), x, y, z)
    (y, z) = (rng.randn(2), rng.randn(2))

    def loop(loop_impl, x):
        return loop_impl(cond, body, (0, x, y, z))[1]
    loop_lax = partial(loop, lax.while_loop)
    loop_ref = partial(loop, while_loop_reference)
    x = rng.randn(2)
    ans = api.jvp(loop_lax, (x,), (x,))
    expected = api.jvp(loop_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(loop_lax, (x,), order=2, modes=['fwd'])
def testWhileJVPWithGrowingNonzeroTangents(self):
    rng = np.random.RandomState(0)

    def cond(state):
        (i, x, y, z) = state
        return (i < 2)

    def body(state):
        (i, x, y, z) = state
        y = (x * x)
        z = (y * y)
        return ((i + 1), x, y, z)
    (y, z) = (rng.randn(2), rng.randn(2))

    def loop(loop_impl, x):
        return loop_impl(cond, body, (0, x, y, z))[1]
    loop_lax = partial(loop, lax.while_loop)
    loop_ref = partial(loop, while_loop_reference)
    x = rng.randn(2)
    ans = api.jvp(loop_lax, (x,), (x,))
    expected = api.jvp(loop_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(loop_lax, (x,), order=2, modes=['fwd'])