Optimization with state

Suppose I have an optimization problem where the function to be optimized requires solving some expensive iterative sub-problem. It would be efficient if I could re-use the solution from the last step of the optimizer to warm start the next step. Something like

def myfun(x, y0):
    y = expensive_subproblem(x, initial_guess=y0)
    return some_other_func(x,y), y

And a sketch of what I'm doing now:

for _ in range(maxiter):
    ((f, y), df) = jax.value_and_grad(myfun, argnums=0, has_aux=True)(x, y)
    x = optimizer_step(x, f, df)

Is something like this possible in optimistix?

No repository files or tests are named. Start by reading the optimizer API and tracing how objective values, gradients, and auxiliary state move through an optimization step. Done would require a decided approach for carrying warm-start state between iterations, with documented behavior and tests for the stateful objective case.