Hello, PyMC developer here. We're excited to give Theano a second life, and hope that this work will lend some staying power to the PyMC project in the probabilistic programming world.
As always, we're happy to accept to contributions! If you're looking to get involved, now is a great time. Please don't hesitate to speak up or reach out, either on the Theano-PyMC GitHub repo (https://github.com/pymc-devs/Theano-PyMC) or some other way (my website's in my bio).
Good to hear Theano has come back to life, it was my first deep learning library :). These days I do PyTorch, and while I really appreciate its debuggability and flexibility, I definitely can see in retrospect some of the advantages of Theano's declarative approach.
In that regard, I am curious about why Tensorflow didn't work out. I understand Tensorflow version 1 implements a declarative mode that I guess is in many ways similar to Theano's. I'm assuming v2 still supports that mode, on top of the new eager mode -- is that the case? If so, was there some aspect of its implementation that made it unsuitable for PyMC?
I'm only one person involved, but my primary reason for choosing Theano over TensorFlow has to do with the ability to manipulate and reason symbolically about models/graphs.
In order to improve the performance and usability of PyMC, I believe we need to automate things at the graph level, and Theano is by far the most suitable for this--between the two, at least.
You can find some work along these lines in the Symbolic PyMC project (https://github.com/pymc-devs/symbolic-pymc); it contains symbolic work done in both Theano and TensorFlow.
Really excited by this as well. I couldn’t really wrap my head around the proposed API changes for PyMC4. I find PyMC3 to be a pleasure to work with and understand!
Will the JAX backend and integration with external JAX modules mean that we’ll see improvements to PyMC3’s variational inference module? That would really increase the versatility of PyMC3 for probabilistic modelling in Python.
I wonder if this was borne out of the success of NumPyro. Been playing around with, pystan and NumPyro recently. Mostly because pymc3 is going the way of the dodo.
NumPyro / JAX / PyTorch just seems like the most versatile offering out there right now
I'm the author of the JAX backend in Theano, and, no, this didn't have anything to do with NumPyro--especially since I neither use nor know very much about NumPyro.
It was just a quick demonstration of how easily one can use Theano as a generalized graph "front-end", while also preserving its more unique and programmable symbolic optimization capabilities. JAX was one of a few "backends" I considered, and, due to the JAX Python library, it also looked like the most straightforward one to implement first.
on the contrary, i recently was looking for python libraries for some bayesian computations for some greenfield development and pymc3 was at the top of the list. with statistical libraries, i prioritize well-tested, large community, and extensive documentation. if others share the same priorities, there's a long future for pymc3.
That's good to know. My comment was borne out of the uncertainty around pymc4 with tensorflow. We are currently using Stan in production and have been hesitant to pull the trigger on pymc3/4 due to this uncertainty.
I welcome and applaud the choice of JAX as it shows a lot of promise with autograd and flexible execution targets.
I have a Jax question here - could someone explain what the PyMC developers exactly did to move their C/C++ backend to Jax ?
Because this sounds very numba-esque. I always thought Jax was just a math library, that was slightly more usable than Numpy. This article makes it seem that code written in Jax ends up being significantly faster than Numpy...almost close to C
JAX can be used as a drop-in replacement for numpy but it enables JIT compilation (similarly to numba) and, more important, very efficient GPU compilation.
Numba is great for CPU-only code, JAX would be my goto solution to add GPU on top of a numpy code.
Theano had a numpy backend so moving to JAX was very straightforward and enabled GPU computations which is now a must for fast deep learning.
JAX does JIT compilation which is not done by Cupy so you get better performance by fusing operations and compiling them cleverly ([0] is a good benchmark to get a feeling for the performance of the various options).
However, it is a trade-off, the jitting time will increase your run time (you can expect at least 10 seconds of compiling time) and JAX does not let you do in-place operations (however, it has workarounds).
If you can afford the time needed to jit your code then I highly recommend JAX.
Related question. I'm trying to sample something with no gradient information (messy astronomical likelihood computation). I'm currently using emcee which uses a Goodman-Weare sampler which breaks down with too many parameters. Are there any better samplers which don't need a gradient?
Maybe variational inference is a possibility? You can try searching for black box stochastic gradient variational inference, easy to use and scales well.
As always, we're happy to accept to contributions! If you're looking to get involved, now is a great time. Please don't hesitate to speak up or reach out, either on the Theano-PyMC GitHub repo (https://github.com/pymc-devs/Theano-PyMC) or some other way (my website's in my bio).