Theano, TensorFlow and the Future of PyMC
PyMC3 is an open-source library for Bayesian statistical modeling and inference in Python, implementing gradient-based Markov chain Monte Carlo, variational inference, and other approximation methods. These algorithms currently rely on Theano for computation, specifically for providing gradients.
Posted by Chris Fonnesbeck on behalf of the PyMC development team
Since the Theano team announced that it would cease development and maintenance of Theano within a year, we, the PyMC developers, have been actively discussing what to do about this. In this post we want to make two big announcements:
PyMC4 will be built on TensorFlow Probability
We are very excited to announce that the new version of PyMC will use TensorFlow Probability (TFP) as its backend. TensorFlow already has a very broad user base and with TFP gained a powerful new library with elegant support for probability distributions and transformations (called bijections, see the TFP paper for a full description), as well as a layer for constructing probabilistic models, called Edward2.
It is clear that TFP’s focus is to provide a strong foundation upon which flexible statistical models for inference and prediction can be constructed from the ground up. Its focus is not, however, to provide a high-level API which makes construction and fitting of common classes of models easy for applied users. Instead, this has always been PyMC’s core mission.
So, in our deliberations regarding the future of PyMC, we were delighted when the TFP developer team reached out to invite us to collaborate on making PyMC4 a high-level interface for TFP that focuses on usability and high-level model specification and inference. We will aim to support the same variety of models currently supported by PyMC3, but allowing for the construction of custom models using the TFP/Edward2 abstraction layer when required. It has always been a dream of ours to furnish users with a Magical Inference Button™ with which PyMC would automatically choose the optimal parameterization and fitting method for their model.
Leveraging TFP has many benefits for our developers and users alike. TensorFlow is widely supported and is a de facto industry standard, making it a viable long-term solution for supporting computation for PyMC. With TFP, we can stand on the shoulders of giants and benefit from the work being done to build a highly performant and scalable statistical package with out-of-the-box support for GPUs and TPUs, as well as state-of-the-art inference algorithms including MCMC and variational inference.
We have already begun taking steps to make PyMC4 a reality: we have recruited a Google Summer of Code student, Sharan Yalburgi, who will be building a prototype using TensorFlow, while one of our core developers, Osvaldo Martin, has lead the charge in splitting out PyMC3’s output visualization and analysis code into a new library called ArviZ, which will ease portability. Realistically, however, the road to a stable and performant PyMC4 will be a long one, so we do not expect a production-ready replacement for PyMC3 to be ready any time soon.
This brings us to our second announcement…
The PyMC team will take over Theano maintenance
While discussing the pros and cons of candidate computational backends, including TensorFlow, PyTorch, and MXNet,we realized that Theano remains a fantastic library that already does exactly what we want it to do: it is mature, fast, well packaged and has continues to serve PyMC3 well. We are extremely proud of PyMC3and its capabilities under Theano, and know that it is used in production systems and research labs around the world. We care deeply about our users and want you to feel safe in continuing to use it and build on it.
To this end, we have been in discussions with Frederic Bastien (Theano’s maintainer and BDFL) regarding assuming control of Theano maintenance once they step down. Over the next year, the PyMC devs will get access to the Theano infrastructure (GitHub, PyPI, etc.) and become better acquainted with the process for PRs and releases.
Critically, we do not plan to run Theano as Frederic did: we will be supporting Theano to the extent that it serves as PyMC3’s computational engine. As such, we will not add features, provide support, or otherwise improve Theano for issues not related directly to PyMC3. This is due entirely to following our mission of producing great probabilistic programming software, which does not permit us to devote time and resources to Theano beyond its role under PyMC3. Given the maturity of the code base, we are confident of a smooth transition and anticipate the effort of ensuring that Theano is still compatible with updates to its dependencies will be straightforward.
This commitment has three key benefits:
- New and existing users have certainty that PyMC3 development and support will continue
- It may inspire others to get involved with keeping Theano afloat (please get in touch if that’s you)
- It gives us more time to develop PyMC4 more deliberately
Its important to remember that PyMC3 is still our flagship project, and we will continue to develop it as such (see our roadmap). New features continue to be added, bugs are diligently being fixed, and we are sticking to a regular release schedule over the coming years. Far from being a distraction, working on PyMC4 and exploring the capabilities of different backends has been a source of inspiration for improving PyMC3, even in fundamental ways (e.g, WIP PR to refactor distributions). Thus, PyMC3 will continue to be the primary target of our development efforts and we aim to prepare a long term support release once PyMC4 is in beta.
Moving Forward
We would like to thank every user and contributor to PyMC over the years. It has been very rewarding to see both the community and the core developer team grow by leaps and bounds over the last 2 years, and we are excited for what the future holds as we embark on this next chapter.
© 2018 the PyMC development team
