def body(i, xy):
    (x, y) = xy
    y = (y + jax.grad((lambda z: jnp.sum(jnp.maximum(z, 0.0))))(x))
    return (x, y)
def body(i, xy):
    (x, y) = xy
    y = (y + jax.grad((lambda z: jnp.sum(jnp.maximum(z, 0.0))))(x))
    return (x, y)