def testIssue757(self):

    def fn(a):
        return jnp.cos(a)

    def loop(val):
        iterations = 10

        def apply_carry(x, i):
            return (api.grad(fn, argnums=(0,))(x)[0], i)
        (final_val, _) = lax.scan(apply_carry, val, jnp.arange(iterations))
        return final_val
    arg = 0.5
    api.jit(api.jacfwd(loop, argnums=(0,)))(arg)
def testIssue757(self):

    def fn(a):
        return jnp.cos(a)

    def loop(val):
        iterations = 10

        def apply_carry(x, i):
            return (api.grad(fn, argnums=(0,))(x)[0], i)
        (final_val, _) = lax.scan(apply_carry, val, jnp.arange(iterations))
        return final_val
    arg = 0.5
    api.jit(api.jacfwd(loop, argnums=(0,)))(arg)