How to run scipy's BFGS on GPU

Question

I'd like to run scipy implementation of BFGS optimization algorithm on GPU and scipy seems not to support GPUs. The target function which I want to run on GPU is the following one which is part of the implementation of this repository:

//here the variable initializations

opts = {'gtol': effective_gamma, 'norm': 2}
result = minimize(private_loss, x0, (x, y), method='BFGS', jac=private_gradient, options=opts,
callback=cb)

I know there is Tensorflow Probablity implementation of BFGS, but I couldn't find out how I can convert this scipy function into Tensordlow Probablity. Any Idea how I could to run the following function on GPU with minimum code change?

Answer 1

Here are two suggestions of mine:

• jax.scipy: Jax contains the implementation of scipy and also supports GPUs and TPUs. You can theoretically install it, convert your numpy variables to jax.numpy and call jax.scipy.optimize.minimize(params):

   import jax.numpy as jnp
   from jax.scipy.optimize import minimize
  
   // here x0,x,y initialization and private_loss and private_gradient functions definition 
  
   x0 = jax.numpy.asarray(x0)
   x = jax.numpy.asarray(x)
   y = jax.numpy.asarray(y)
  
   opts = {'gtol': effective_gamma, 'norm': 2}
   result = minimize(private_loss, x0, (x, y), method='BFGS', jac=private_gradient, options=opts, callback=cb)  
  
   // here rest of the code
  

Dont foget to also take care of converting the variables which will be used in private_loss and private_gradient function to jax.numpy.

• Tensorflow Probability: As you already mentioned, you can also use bfgs_minimize implementation by tensorflow. based on the colab here, your code will be something like this:

   def make_val_and_grad_fn(value_fn):
     @functools.wraps(value_fn)
     def val_and_grad(param):
       return tfp.math.value_and_gradient(value_fn, param)
     return val_and_grad
  
   def run(optimizer):
     result2 = optimizer()
     return np_value(result2)
  
   @make_val_and_grad_fn
   def private_loss(param):
       // here private_loss
  
   x_var = tf.Variable(x)
   y_var = tf.Variable(y, dtype = tf.float64)
   x0 = tf.Variable(x0)
  
   tolerance = effective_gamma
  
   @tf.function
   def minimize_with_bfgs():
       result1 = tfp.optimizer.bfgs_minimize(
       private_loss,
       initial_position=tf.constant(x0),
       tolerance=tolerance)
       return result1
  
   results = run(minimize_with_bfgs)