Unnecessary recompilation of _while_loop_lax

ed1febef279de5e1120af70dff04261b61175383 claims that jitting `_while_loop_lax` is redundant. However, this change also seems to prevent the result from being cached, causing recompilation of loops like https://github.com/google/jaxopt/blob/58bac0ac375732dce358cf85583ae7fe3632b8cf/jaxopt/_src/base.py#L314 Reverting ed1febef279de5e1120af70dff04261b61175383 drastically reduces compilation times for my use case, so it probably makes sense to address this cache issue.