JIT function frozen... minimization failure

[jecampagne]

Hi,
Trying to make a helper function to allow for jit-method in class, I face a problem in the following use-case of minimization when I would like to minimize the number of recompilation of the _test function. For the sake of completeness I use the jaxopt.ScipyBoundedMinimize minimizer which cannot be jitted.

Below is the snippet where for "res 1" I instantiate a new A object at each time which yields a recompilation at each time and gives the right expected solution (minimum is equal to a). For the "res 2" case, a single A object is instantiate once for all, and it yields a single compilation of _testBUT the results are only correct for the first avalue.

I do not yet manage to reconcile 1) right solution and 2) a single compilation. Have you a comment/idea? Thanks in advance.

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

import jaxopt

@jit
def _test(x, a):
    print("compile....")
    return (x-a)**2

class A():
    def __init__(self):
        
        tol = 1e-3
        method = 'L-BFGS-B'
        options = {'disp':False,'ftol':tol, 'gtol':tol, 'maxiter':600}
        self.jscMin=jaxopt.ScipyBoundedMinimize(fun=self.test,
                                           method=method, 
                                           tol=tol,
                                           options=options)
    
        self.Init()
    
    def Init(self):
        self.a = None
        
    def set_a(self, val):
        self.a = val
    
    def test(self, x):
        return _test(x, self.a)

    def optimize(self, init_val, low_bnd, high_bnd):
        res = self.jscMin.run(init_val,bounds=(low_bnd,high_bnd))
        return res
    
        
    

print(">>> res1")
for a in [2., -5., 10.]:
    obja = A()
    obja.set_a(a)
    print("a: ",obja.a)
    res = obja.optimize(jnp.array(0.),-10.,10.)
    print(res)

    
print(">>> res2")
obja = A()
for a in [2., -5., 10.]:
    obja.set_a(a)
    print("a: ",obja.a)
    res = obja.optimize(jnp.array(0.),-10.,10.)
    print(res)

leading to

>>> res1
a:  2.0
compile....
OptStep(params=DeviceArray(2., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
a:  -5.0
compile....
OptStep(params=DeviceArray(-5., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
a:  10.0
compile....
OptStep(params=DeviceArray(10., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=1))


>>> res2
a:  2.0
compile....
OptStep(params=DeviceArray(2., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
a:  -5.0
OptStep(params=DeviceArray(2., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
a:  10.0
OptStep(params=DeviceArray(2., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))

[jecampagne]

Ha, I got it

In fact, jaxopt.ScipyBoundedMinimize implicitly trigger a compilation to minimize. So in my original snippet, this is the test(self,x). But then, comes the self problem which is frozen, so the set_a is completely useless as the first avalue known at compilation is 2.

In the code below, the function argument of jaxopt.ScipyBoundedMinimize is directly the helper function, and I adjust the argument of self.jscMin.run in optimize to set the proper a-value

@jit
def _test(x, a):
    print("compile....")
    return (x-a)**2

class A():
    def __init__(self):
        
        tol = 1e-3
        method = 'L-BFGS-B'
        options = {'disp':False,'ftol':tol, 'gtol':tol, 'maxiter':600}
        self.jscMin=jaxopt.ScipyBoundedMinimize(fun=_test,
                                           method=method, 
                                           tol=tol,
                                           options=options)
    
        self.Init()
    
    def Init(self):
        self.a = None
        
    def set_a(self, val):
        self.a = val
    
#    def test(self, x):
#        return _test(x, self.a)

    def optimize(self, init_val, low_bnd, high_bnd):
        res = self.jscMin.run(init_val,bounds=(low_bnd,high_bnd), a=self.a)
        return res
 

It yields

>>> res2
a:  2.0
compile....
OptStep(params=DeviceArray(2., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
a:  -5.0
OptStep(params=DeviceArray(-5., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=2))
a:  10.0
OptStep(params=DeviceArray(10., dtype=float64), state=ScipyMinimizeInfo(fun_val=DeviceArray(0., dtype=float64, weak_type=True), success=True, status=0, iter_num=1))

=> 1 single compilation and all minimisations are correct !. It is a good solution or do you see some possible failures?

[jecampagne]

Well, optax.adam is absolutely not adapted :( . I will try jaxopt.ProjectedGradient with .projection_box, but do not know it is JAX-compatible (contrary to ScipyBoundedMinimize)

[jecampagne]

Well after investigation jaxopt.ProjectedGradient with .projection_box is not adapted too. So I need certainly to find a L-BFGS-B solver with bounded hyper-cube constraint compatible with JAX...

[PhilipVinc]

how comes projected gradient is not jax compatible?

[jecampagne]

ho, may be my english is not ad equa, jaxopt.ProjectedGradient is certainly JAX compatible, but for my use case it does not manage to get right.

[jecampagne]

I just see that jax.scipy.optimize.minimize has a BFGS method, (and probably LBFGS is in the pipeline) but does it is possible to impose "box-bounds" on parameters? Thanks