def apply_carry(carry, _):
    (i, x) = carry
    new_x = (x - (0.1 * api.grad(energy_fn)(x)))
    new_carry = ((i + 1), new_x)
    return (new_carry, _)
def apply_carry(carry, _):
    (i, x) = carry
    new_x = (x - (0.1 * api.grad(energy_fn)(x)))
    new_carry = ((i + 1), new_x)
    return (new_carry, _)