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
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