ScipyBoundedMinimize : recompilation of function to minimize at each instantiation?

[jecampagne]

Hello,

Here a simple snippet:

import jax
import jax.numpy as jnp
from jax import vmap, jit
import numpy as np
jax.config.update("jax_enable_x64", True)

import jaxopt

@jit
def test(x):
    print("compile....")
    return -(x+2.)**2

tol = 1e-3
method = 'L-BFGS-B' # 'SLSQP'
options = {'disp':False,'ftol':tol, 'gtol':tol, 'maxiter':600}
jscMin=jaxopt.ScipyBoundedMinimize(fun=test,
                                           method=method, 
                                           tol=tol,
                                           options=options)

res = jscMin.run(jnp.array(10.),bounds=(-100,100))
print(">>> res1", res)

jscMin=jaxopt.ScipyBoundedMinimize(fun=test,
                                           method=method, 
                                           tol=tol,
                                           options=options)



res = jscMin.run(jnp.array(5.),bounds=(-100,100))
print(">>> res2", res)



res = jscMin.run(jnp.array(-5.),bounds=(-100,100))
print(">>> res3", res)

which leads to

compile....
>>> res1 OptStep(params=DeviceArray(100., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(-10404., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
compile....
>>> res2 OptStep(params=DeviceArray(100., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(-10404., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
>>> res3 OptStep(params=DeviceArray(-100., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(-9604., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))

Is it foreseen that there is a compilation of test for each instantiation of ScipyBoundedMinimize ? that is to say to avoid recompilation then the user should use a single jscMin ?

The use-case is that, one needs sometimes to use different initial start to proceed to a choice of best parameters, so calling several times ScipyBoundedMinimize could have an overhead simply due to recompilation of the user function to minimize (here test but could be dramatically more complicated).

Any hint to optimize better ? Thanks.