Hot fix: corrected condition in lbfgs

[zaccharieramzi]

The feature I had introduced in #323 was failing when the run function was jitted and was a no-op when not because of the following reason:

 ~True == -2  # this is True

Therefore when jitted it was complaining about different types in a condition function, and when not jitted it was equivalent to always being False.

EDIT

Actually I am still running into an error when jitted, so will continue to investigate.

The gist of the error is Abstract tracer value encountered where concrete value is expected, basically doing (not self.stop_if_linesearch_fails | ~state.failed_linesearch) is not allowed because one is a bool and the other is an abstract value.

[zaccharieramzi]

With a jnp.logical_or the second issue was resolved

[mblondel]

Could you add a non-regression test? Thanks

[zaccharieramzi]

Yes just did that: my diagnosis was incorrect it was not jitting the run that was problematic but rather not using implicit diff (which I am not in my use case).
Therefore I just added a case in the rosenbrock test, which fails when I don't have the current fix, and doesn't now.

[mblondel]

Actually I am still running into an error when jitted, so will continue to investigate.

Is this one fixed too?

[zaccharieramzi]

Yes actually the problem was not coming from it being jitted but from setting implicit_diff to False, I had the wrong diagnosis

[mblondel]

Perfect! :) Please squash and I'll merge.

[zaccharieramzi]

Done you can merge when all green in CI