How to run scipy's BFGS on GPU

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?

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)