Fix complex support for L-BFGS

[gautierronan]

Closes #1141.

Not 100% sure that this doesn't break other things or fully works, but at least the MWE below seems to work fine.

import optax
import jax.numpy as jnp

def f(x):
    return jnp.sum(jnp.abs(x**2))

solver = optax.lbfgs()
params = jnp.array([1.0 + 1.0j, 2.0 + 2.0j, 3.0 + 3.0j])
print("Objective function: ", f(params))

opt_state = solver.init(params)
value_and_grad = optax.value_and_grad_from_state(f)

for _ in range(5):
    value, grad = value_and_grad(params, state=opt_state)
    updates, opt_state = solver.update(
        jnp.conj(grad), opt_state, params, value=value, grad=jnp.conj(grad), value_fn=f
    )
    params = optax.apply_updates(params, updates)
    print("Objective function: ", f(params))

Notice the solve.update call which requires a jnp.conj(grad) twice. I believe this is correct and aligned with other optax solvers, but not sure either.

[vroulet]

Hey @gautierronan,
Thanks for the PR! We'll need a test. Take look at this PR: google/jaxopt#468 that added support for complex parameters for the lbfgs of jaxopt. I think you'll find all that you'll need in that PR.
Thanks again!

[gautierronan]

@vroulet Should be good for review. Note that, in the test, I have commented one linesearch option because the test was not passing, but I don't think it's related to complex support.

[vroulet]

That looks great, thank you @gautierronan ! I left you some cosmetic comments.
In addition, can you add the following note in the docstrings of lbfgs, scale_by_zoom_linesearch, scale_by_backtracking_linesearch (please compile the docs to verify that it's well formatted)?

  ..note:: The algorithm can support complex inputs.

Ideally, one would add a simple example in the lbfgs notebook to link this note to an implementation (as you pointed out the gradients need to pass through a conjugate). Let me know if you could add that too.

Also could you add the test on the complex rosenbrock that was in jaxopt. It should amount to add exactly

  def test_lbfgs_complex_rosenbrock(self):
    # Taken from previous jax tests
    tol = 1e-5
    complex_dim = 5

    fun_real = _get_problem('rosenbrock')['fun']
    init_real = jnp.zeros((2 * complex_dim,), dtype=complex)
    expected_real = jnp.ones((2 * complex_dim,), dtype=complex)

    def fun(z):
      x_real = jnp.concatenate([jnp.real(z), jnp.imag(z)])
      return fun_real(x_real)

    init = init_real[:complex_dim] + 1.j * init_real[complex_dim:]
    expected = expected_real[:complex_dim] + 1.j * expected_real[complex_dim:]

    opt = alias.lbfgs()
    got, _ = _run_opt(opt, fun, init, maxiter=500, tol=tol)
    chex.assert_trees_all_close(got, expected)

If it does not pass, let me know. I can work on that in a separate PR once yours is merged.

[gautierronan]

Thanks for the review. All comments should be addressed in the last commits. Let me know if that works.

Ideally, one would add a simple example in the lbfgs notebook to link this note to an implementation (as you pointed out the gradients need to pass through a conjugate). Let me know if you could add that too.

I don't think I will have time for this. Also, it's IMO expected that gradients need to pass through a conjugate (nothing specific to L-BFGS), and is already the behavior in the other optax functions I believe. So a tutorial dedicated to complex-valued gradient descent might be more relevant than a specific addition to the L-BFGS tutorial?

[vroulet]

So a tutorial dedicated to complex-valued gradient descent might be more relevant than a specific addition to the L-BFGS tutorial?

Good point. We can make an issue for that. Thanks for all!