I’ve checked your proposal. I still suggest not taking OPVI way for PyMC4 GSoC. There are a couple of reasons for me to be skeptical about it:
- Implementation was hard and took around a year of continuous development and achieving speed parity with its previous functional implementation (for a very long time we had legacy pm.advi because of performance reasons). For this reason, I see OPVI way to be a dead way. I hardly believe it would be doable in 3 months.
- Given point one is not a concern and PyMC3 OPVI implementation is cristal clear, TF2 lacks core abstractions I relied on during my work on OPVI. Namely, symbolic graph representation,
theano.clone, (easy) vectorized input to graph. Without these tools/properties it would be challenging to “just implement” the concept.
- Given the above points are still not a concern, what I doubt, there is yet no need to put a lot of abstractions for the inner implementation. You may see, OPVI in PyMC3 is working more like building blocks for VI (KL, Stein, Amortized). We yet have no blocks in PyMC4. That is why there is still no need for that level of abstractions.
As for PyMC4 project, we need to create a simple API, similar to
pm.fit to provide few VI implementations, such as MeanField/FullRank ADVI. So far, the top abstraction level is an Approximation. And still, TF2 has its caveats, VI implementation will require a sort of graph manipulation which we only can access in a too limited way. I see this level as minimal required to implement VI, but modifications to pm4 Executor and its logic are very likely
I hope we’ll get the need in OPVI, but not in the nearest future. Anyway, Approximation is a part of OPVI and the future work may start out of this milestone (that is a GSoC project).
I suggest rewriting the proposal narrowing the scope for MeanField and FullRank ADVI only.