Efficient NewtonCG Implementation about optimistix HOT 3 OPEN

Hi all, thanks for the phenomenal library. We're already using it in several statistical genetics methods in my group!

That's great to hear, thank you!

On HVPs: if I understand you correctly, this is a general JAX question, rather than specifically a question of how to integrate a solve into Optimistix? You're looking to get both the gradient and a HVP without having to treat them both separately (which would be 3 sweeps in total).
This approach will get you 2 sweeps, which I think is optimal:

def to_jvp(x):
    return jax.value_and_grad(fn)(x)

(f, dfdxi), (_, dfdxidxjvj) = jax.jvp(to_jvp, (x,), (v,))

from optimistix.

Newton CG is one of the algorithms which I think could be somewhat involved to implement. It is further from the existing solvers in Optimistix, so there's more custom work that needs to be done to get it running.

There's two steps I would take if I were implementing it (which I may in the future):

1. Implement a custom top-level solver which passes the Hessian-vector product $h(v) = \nabla^2 f(x)$ via lx.FunctionLinearOperator in FunctionInfo.EvalGradHessian. Take a look at the abstract solvers (AbstractGradientDescent, AbstractGaussNewton, AbstractBFGS, etc.) for a bit on how to do this, and check out the docs for FunctionInfo. Implementing a top-level solver from scratch for the first time can be a little tricky, because it requires touching all of the abstractions used in Optimistix. However, the technical bit should be very simple here. If you find it difficult to figure out from looking at existing implementations just poke me!
2. Implement an early-exit version of CG in Lineax. This is needed because for inexact Newton CG the CG iterates stop immediately if a direction of negative curvature is encountered, allowing non-psd matrices to be approximately solved with CG. This should be straightforward, copy the code from the implementation of CG in Lineax and add the curvature condition to the function not_converged (something like tree_dot(y, operator.mv(y)) > 0. This may need tweaking, and you'll also have to disable the positive semidefinite checks.)

This should be pretty much everything though. Just use optx.NewtonDescent with linear_solver=NewEarlyExitCG for the descent, and your favorite Optimistix Search for the search and it should all work! You could probably use some other descents as well, such as DoglegDescent or DampedNewtonDescent by passing the new linear solver in to these descents. I don't know how well these would work, as they're not standard algorithms (neat!)

from optimistix.

Yes exactly! It isn't clear to me how to work this into a Newton-like CG solver, but I'll keep toying around. Thanks as always for your enthusiasm and detailed help. It is greatly appreciated!

from optimistix.