@parameterized.named_parameters((dict(testcase_name='_loop={}'.format(loop), loop=loop) for loop in ['while', 'fori', 'fori_inside_cond', 'fori_inside_scan']))
def testWhileGradError(self, loop: str='fori_inside_scan'):
    if (loop == 'while'):
        func = (lambda x: lax.while_loop((lambda i: (i < 5.0)), (lambda i: (i + 1.0)), x))
    elif (loop == 'fori'):
        func = (lambda x: lax.fori_loop(x, (x + 2.0), (lambda i, c: c), x))
    elif (loop == 'fori_inside_jit'):
        func = api.jit((lambda x: lax.fori_loop(x, (x + 2.0), (lambda i, c: c), x)))
    elif (loop == 'fori_inside_cond'):
        func = (lambda x: lax.cond(True, x, (lambda x: lax.fori_loop(x, (x + 2.0), (lambda i, c: c), x)), 1.0, (lambda x: x)))
    elif (loop == 'fori_inside_scan'):
        func = (lambda x: lax.scan((lambda c, x: (lax.fori_loop(x, (x + 2.0), (lambda i, c1: (c1 * c)), x), None)), x, np.ones(2))[0])
    else:
        assert False
    with self.assertRaisesRegex(ValueError, 'Reverse-mode differentiation does not work for lax.while_loop'):
        api.grad(func)(1.0)
    api.linearize(func, 1.0)
@parameterized.named_parameters((dict(testcase_name='_loop={}'.format(loop), loop=loop) for loop in ['while', 'fori', 'fori_inside_cond', 'fori_inside_scan']))
def testWhileGradError(self, loop: str='fori_inside_scan'):
    if (loop == 'while'):
        func = (lambda x: lax.while_loop((lambda i: (i < 5.0)), (lambda i: (i + 1.0)), x))
    elif (loop == 'fori'):
        func = (lambda x: lax.fori_loop(x, (x + 2.0), (lambda i, c: c), x))
    elif (loop == 'fori_inside_jit'):
        func = api.jit((lambda x: lax.fori_loop(x, (x + 2.0), (lambda i, c: c), x)))
    elif (loop == 'fori_inside_cond'):
        func = (lambda x: lax.cond(True, x, (lambda x: lax.fori_loop(x, (x + 2.0), (lambda i, c: c), x)), 1.0, (lambda x: x)))
    elif (loop == 'fori_inside_scan'):
        func = (lambda x: lax.scan((lambda c, x: (lax.fori_loop(x, (x + 2.0), (lambda i, c1: (c1 * c)), x), None)), x, np.ones(2))[0])
    else:
        assert False
    with self.assertRaisesRegex(ValueError, 'Reverse-mode differentiation does not work for lax.while_loop'):
        api.grad(func)(1.0)
    api.linearize(func, 1.0)