Hager-Zhang linesearch does not work with non-jittable functions

Currently, the Hager-Zhang linesearch uses its own `lax.while_loop` (see [this line](https://github.com/google/jaxopt/blob/main/jaxopt/_src/hager_zhang_linesearch.py#L200)). This means that the objective functions must be jittable.

In order to support non-jittable functions, we should use JAXopt's while_loop instead, so that we can propagate the jit and unroll options to it.

CC @emilyfertig @srvasude

See #444 for context