OptaxSolver how to proceed? about jaxopt HOT 6 CLOSED

The optimizers of Optax are meant for minimization, not root finding. If we take a look at your function neg_Lag we can see that the minimum does not exist: the Lagrange multiplier x[3] is allowed to take any value in , so as long as your surface is not zero it can be used to reach any real number. Optax is working well by diverging since the minimum does not exist anyway.

So you must:

• either use an optimizer to perform constrained optimization of your function vol; here ProjectedGradient might be deceiving since your feasible set is not convex.
• either use a root finding algorithm to find the zero of jax.grad(Lag). Unfortunately, currently, we lack options for multidimensional root finding (best we have currently is ScipyRootFinding). Other options based on Fixed point finding will be availble soon.
• reformulate the multidimensional root finding algorithm into an optimization problem: instead of finding p such that we seek to minimize

The latter works:

opt = optax.adagrad(0.1)

@jax.jit
def objective_fun(p):
  delta = gLag(p)
  return jnp.sum(delta**2)  # minimize gradient norm

solver = jaxopt.OptaxSolver(opt=opt, fun=objective_fun, maxiter=2000)
init_params = jnp.array([1.5,0.5,1.0,0.1])
params, state = solver.init(init_params)
print('init', params, neg_Lag(params))

@jax.jit
def jitted_update(params, state):
  return solver.update(params=params, state=state)
for i in range(20*1000):
    params, state = jitted_update(params, state)
    if i%100 == 0: 
        print(i, params, Lag(params), objective_fun(params))

However this algorithm is far from being efficient.

from jaxopt.

@mblondel you are absolutely right, my problem is a root finding

and not a minimization. Sorry I have forget this point when I was using my (old) solveLagrangian that I have jaxized , it is exactly doing that root search thanks to Newton step.

Thanks for your different method discussion and snipped, too. I am not sure that I can contribute but your lib is really nice and I encourage for new code implementation.

from jaxopt.

@mblondel

rf = jaxopt.ScipyRootFinding(optimality_fun=gLag, method='hybr')
rf.run(jnp.array([1.5,0.5,1.0,0.1]))

gives

OptStep(params=DeviceArray([2. , 2. , 2. , 0.5], dtype=float64), state=ScipyRootInfo(fun_val=DeviceArray([-2.05051975e-10,  3.02247116e-11, -3.61966457e-10,
              6.25949070e-10], dtype=float64), success=True, status=1))

So I wander why you reject this method in your comment? may be I have misunderstood something.

from jaxopt.

So I wander why you reject this method in your comment? may be I have misunderstood something.

I don't reject it, I was just mentioning that we had nothing else for this purpose.

For some reason you wanted to use Optax so I showed you an example with Optax.

But ScipyRootFinding is fine too.

from jaxopt.

Note: I am not @Algue-Rythme :)

from jaxopt.

ho sorry @mblondel, I was also in contact with @Algue-Rythme in an other thread :)
Now, your code with Optax was really nice.
Thanks a lot.

from jaxopt.