performing many gradient-based optimizations in parallel with TensorFlow
I believe the BDF optimizer and LBFGS both support batches (of problems), so you could have an outer "problem" axis to your data and leastsq return value. But since BDF is for stiff problems, it's likely to have much longer runtimes for some problems than for others, and you might be best served treating each problem separately (tf.map_fn) as opposed to running them all in parallel -- in a batch, you can't run ahead onto the next LBFGS iteration for problem X until you compute the BDF integration for problem Y. Or just use a python for loop over your problems, each time calling a @tf.function def lbfgs_over_bdf(data): ....
I have a model which requires solving a system of ODEs with tfp.math.ode.BDF, and I would like to find the individual least-squares fits of this model to n > 1000 datasets. That is to say, if my model has m parameters then at the end of the optimization process I will have an n by m tensor of best-fit parameter values.
What would be the best way to perform this optimization in parallel? At this point I’m planning to define an objective function that adds up the n individual sums of square residuals, and then uses tfp.optimizer.lbfgs_minimize to find the best-fit values of the combined n×m parameters.