@parameterized.named_parameters(({'testcase_name': '_jit_loop={}_jit_body={}_jit_cond={}'.format(jit_loop, jit_body, jit_cond), 'jit_loop': jit_loop, 'jit_body': jit_body, 'jit_cond': jit_cond} for jit_loop in [False, True] for jit_body in [False, True] for jit_cond in [False, True]))
def testWhileJVP(self, jit_loop=True, jit_body=False, jit_cond=True):
    cond = (lambda x: (x[(0, 2)] <= 8))
    body = (lambda x: (x * x))
    if jit_cond:
        cond = api.jit(cond)
    if jit_body:
        body = api.jit(body)
    loop = partial(lax.while_loop, cond, body)
    if jit_loop:
        loop = api.jit(loop)
    loop_ref = partial(while_loop_reference, cond, body)
    x = jnp.arange(9.0).reshape((3, 3))
    ans = api.jvp(loop, (x,), (x,))
    expected = api.jvp(loop_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(loop, (x,), order=2, modes=['fwd'])
@parameterized.named_parameters(({'testcase_name': '_jit_loop={}_jit_body={}_jit_cond={}'.format(jit_loop, jit_body, jit_cond), 'jit_loop': jit_loop, 'jit_body': jit_body, 'jit_cond': jit_cond} for jit_loop in [False, True] for jit_body in [False, True] for jit_cond in [False, True]))
def testWhileJVP(self, jit_loop=True, jit_body=False, jit_cond=True):
    cond = (lambda x: (x[(0, 2)] <= 8))
    body = (lambda x: (x * x))
    if jit_cond:
        cond = api.jit(cond)
    if jit_body:
        body = api.jit(body)
    loop = partial(lax.while_loop, cond, body)
    if jit_loop:
        loop = api.jit(loop)
    loop_ref = partial(while_loop_reference, cond, body)
    x = jnp.arange(9.0).reshape((3, 3))
    ans = api.jvp(loop, (x,), (x,))
    expected = api.jvp(loop_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(loop, (x,), order=2, modes=['fwd'])