google / google/jaxopt

Passing a custom callback to jaxopt.ScipyMinimize

Open
#257 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
1.1k
Forks
76
Avg merge
2d 21h
Merged PRs (30d)
1

Description

I would like to ask if there is any way to pass a custom call back function to one of the solvers (TNC) in this case. For example, I am able to make the code below work in scipy.optimize.minimize (BTW, I used jnp in both cases):
```
# This works
res = scipy.optimize.minimize(self.imp_sim_weighted.simulate, ln_par0, \
args = (self.controller.app_data["freq"], self.controller.app_data["z"], lb_col, ub_col, self.smf, self.controller.app_data["weight"]), method = 'TNC',
jac = self.controller.jac_deis, callback= self.imp_sim.callback, options={'maxfun':10000, 'ftol':1e-10, 'xtol':1e-10})
```

But with jaxopts, I get an error:
```
solver = jaxopt.ScipyMinimize(method = "TNC", fun=self.imp_sim.simulate, tol = 1e-12, options ={'maxiter':5000, 'callback':self.imp_sim.callback})
sol = solver.run(ln_par0, self.controller.app_data["freq"], self.controller.app_data["y"], lb_col, ub_col, self.smf, self.controller.app_data["weight"])
```

Error:
```
# res = _minimize_tnc(fun, x0, args, jac, bounds, callback=callback,
# jax._src.traceback_util.UnfilteredStackTrace: TypeError: scipy.optimize._tnc._minimize_tnc() got multiple values for keyword argument # 'callback'
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.