Callback function requires 2 arguments when using method='trust-constr'
- Dominant language
- Python
- Stars
- 1.1k
- Forks
- 76
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 1
Description
This returns an error if a custom callback function is provided, which requires 2 arguments for 'trust-constr'. This
```
solver = jaxopt.ScipyMinimize(fun=self.loss, method='trust-constr', callback=self._metrics,
**kwargs_optimiser)
res = solver.run(params, **extra_kwargs)
```
returns the following error :
```
File [~/Library/Python/3.8/lib/python/site-packages/jaxopt/_src/scipy_wrappers.py:306], in ScipyMinimize._run(self, init_params, bounds, *args, **kwargs)
302 if bounds is not None:
303 bounds = osp.optimize.Bounds(lb=jnp_to_onp(bounds[0], self.dtype),
304 ub=jnp_to_onp(bounds[1], self.dtype))
--> 306 res = osp.optimize.minimize(scipy_fun, jnp_to_onp(init_params, self.dtype),
307 jac=True,
...
--> 431 if callback is not None and callback(np.copy(state.x), state):
432 state.status = 3
433 elif state.optimality < gtol and state.constr_violation < gtol:
TypeError: scipy_callback() takes 1 positional argument but 2 were given
```
I guess this is because _src/scipy_wrappers.py, line 333:
```
if self.callback is not None:
def scipy_callback(x_onp: onp.ndarray):
x_jnp = onp_to_jnp(x_onp)
return self.callback(x_jnp)
else:
scipy_callback = None
```
imposes that the callback function has only one argument.
Contributor guide
Assessment
This issue has not been assessed yet.