PyMC 4.0 Release Announcement

We, the PyMC core development team, are incredibly excited to announce the release of a major rewrite of PyMC (formerly PyMC3): 4.0.0. There are many exciting new updates that we will talk about in this and upcoming blog posts.

It's now called PyMC instead of PyMC3

First, the biggest news: PyMC3 has been renamed to PyMC with this release. PyMC3 version 3.x will stay under the current name to not break production systems. There are a few reasons for this. The main one is that PyMC3 4.0 is quite confusing.

What about PyMC4?

As described in our previous post "The Future of PyMC3, or: Theano is Dead, Long Live Theano", we have experimented with (but given up on) a PyMC4 that was based on TensorFlow and TensorFlow Probability.

One of the major reasons for this was (i) that TF was really difficult to work with, partly because of its eager execution mode, but, more importantly, (ii) we realized that Theano already was in a great spot to take PyMC to the next level.

So to summarize, here are the names and what they refer to:

  • PyMC3 3.x: The version you all love and use, version numbers start with 3.x.
  • PyMC 4.0: The successor to PyMC3, (mostly) API compatible, our main focus going forward, and what this blog post is about. This is the only version that will matter going forward. Version numbers start with 4.x.
  • PyMC4: The discontinued experiment we did with new API on the TensorFlow backend.

Theano -> Theano-PyMC -> Aesara

So after we rediscovered the power of Theano, we went ahead and forked Theano to Theano-PyMC and made PyMC3 rely on this fork instead.

During this forking process we discovered that we could take this project much further still, so we undertook a massive rewrite of Theano-PyMC (this charge was led by Brandon Willard), removing swaths of old and obscure code, and restructuring the entire library to be more developer friendly.

This rewrite motivated a rename of the package to Aesara (Theano's daughter), and is now a project largely separate from PyMC, although PyMC is the main customer.

Most imporantly, however, we added two new computational backends: JAX and numba. The way this works is that aesara is best understood as a computational graph library that allows you to build a computational graph out of array-operations (additions, multiplications, dot-products, indexing, for-loops). From that graph, you can do graph optimizations (log(exp(x)) -> x), symbolic rewrites (N(0, 1) + a -> N(a, 1)), as well as compile that graph, and here we have a ton of flexibility now and can add new compilation backends easily. Currently there is a C, CUDA, JAX (with GPU support), and numba backend. Future blog posts will go into more detail as to how to use these backends and their different trade-offs.

New Features

Most of the changes in v4 happened under the hood and lead to lots of code simplifications and setting the stage for future innovations that were not possible before.

However, there are quite a few good reasons to check out this new version already.

Faster predictive sampling

In v3 we were not really using the full power of Theano. While calling e.g. x = pm.Normal('x') created a theano.TensorVariable which you could manipulate further, like inputting it into another random variable. However, we were not really building a full-fledged graph. This created all kinds of issues but a major one was a very complicated

This is a major change in 4.0.

JAX backend



import aesara
aesara.config.mode = "JAX"  # Default is "Mode"

Numba backend



import aesara
aesara.config.mode = "numba"  # Default is "Mode"

Symbolic shapes

import pymc as pm

with pm.Model() as m:
    x = pm.Poisson('x', 2)
    y = pm.Poisson('y', 2)
    z = pm.Normal('z', shape=(x, y))
    
print(z.eval())
print(z.eval())

Produces

Better NUTS initialization

Scan-support

What's not yet working

Currently we are still missing a few features from v3. These should be hard to add and will be there for the 4.0 stable release. The list includes:

  • Potentials
  • Timeseries distributions
  • Mixture Distributions
Select a repo