Callback function requires 2 arguments when using method='trust-constr'

[martin-millon]

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.