I have an application which uses tf.linalg.lstsq inside of tf.vectorized_map to do least-squares regression to fit parameters for a large set of problem instances.
Similarly, this Bermudan option algorithm calls training to fit a function, inside a loop of the algorithm, for a fixed number of times. It’s the same kind of fitting, just a different choice of fitter.
Generalizing, is it possible to construct a tf.function which does a training iteration on a neural net, which can be called inside tf.vectorized_map? The point again would be to simultaneously do a kind of curve fitting on a large number of similar problem instances in parallel.