The Quickest Migration Guide Ever from PyMC3 to PyMC v4.0

Discover how few changes are needed to upgrade to PyMC v4.0!

Authored by Eric J. Ma on 2022-01-03

PyMC v4.0 is out in beta! PyMC Labs is thrilled to have worked with the PyMC development team to get this release out the door. Usually, a new version implies code changes. In this short post, we’d like to highlight how to migrate your code from PyMC3 to PyMC v4.0.

The tl;dr? While there’s a lot of really cool improvements we’ve made underneath the hood (stay tuned for a full post on this), the number of things you’ll need to change are minimal!

Installation

PyMC v4.0 is now available via pip and conda.

To install by pip:

pip install pymc --pre

(The --pre flag is necessary to install the pre-release before we release the a production-ready version.)

To install by conda:

conda install -c conda-forge pymc
# OR USE MAMBA FOR BLAZING FAST INSTALLATIONS:
mamba install -c conda-forge pymc

Imports

The biggest change you’ll see is probably this one. You probably imported PyMC3 as follows:

import pymc3 as pm  # DO NOT DO THIS!

With PyMC v4.0, you’ll do the following import instead:

import pymc as pm  # DO THIS INSTEAD!

Sampling

In PyMC3, we used to return a MultiTrace object.

with model:
    trace = pm.sample()

In PyMC v4.0, we instead return an ArviZ InferenceData object instead:

with model:
    idata = pm.sample()  # THE NEW CONVENTION!

Underneath the hood, the InferenceData object is a richer data structure than the MultiTrace object, and is essentially a wrapper around xarray. Xarray provides the ability to do much saner indexing and labelling of coordinates. You can think of xarray as “pandas for n-dimensional arrays”. You can read more about the InferenceData object on the official ArviZ docs.

In addition to that the InferenceData object supports storing prior predictive and posterior predictive samples that you can use to sanity-check your models.

with model:
    idata = pm.sample()
    idata.extend(pm.sample_prior_predictive())
    idata.extend(pm.sample_posterior_predictive(idata))

If you do this often, stay tuned for API updates to PyMC v4.0: we are discussing making prior and posterior predictive samples a standard part of pm.sample() itself!

ArviZ

In the past, you may have used PyMC3’s built-in visualization capabilities:

pm.plot_trace(trace)
pm.plot_forest(trace)

Since PyMC3, we’ve started delegating these visualization capabilities to ArviZ, and as of PyMC v4.0, it’s considered idiomatic to just use ArviZ instead.

import arviz as az
az.plot_trace(idata)
az.plot_forest(idata)
az.plot_posterior(idata)

ArviZ has a rich library of Good Bayesian Workflow-supporting visual diagnostics, and we’d encourage you to check it out!

What about all of those custom things?

Perhaps you’ve made something custom using the old Theano library, and are unsure about what to do there?

import theano.tensor as tt

Well, Theano has been replaced by Aesara, and it’s basically a one-to-one replacement (with minor exceptions).

import aesara.tensor as at

For those who are curious, in one Greek tradition, Aesara is the daughter of Theano and Pythagoras.

Aesara and Theano are both symbolic tensor libraries. Since MILA’s discontinued efforts developing Theano, the PyMC dev team took over Theano and, with the upgrades made thus far, we’ve repackaged it into Aesara. You can check out the Aesara docs and GitHub organization to learn more about PyMC’s backing library. While previously Theano was backed by C++, Aesara adds Numba and JAX backends while striving for 100% NumPy API compatibility. It’s an exciting time to be a Pythonista!

And if there’s something that’s really custom that you need, such as a custom JAX-backed operation or high performance custom code in Aesara, then PyMC Labs is here to help!

Summary

In summary, for the vast majority of users of PyMC3, migrating from PyMC3 to PyMC v4.0 is a matter of switching a few lines of code. A few find-and-replace operations throughout your repository, and you should be fine. For the more advanced users, you can use Aesara like Theano. And for those really technically challenging yet high business value problems, PyMC Labs is here for you — just get in touch with us 🙂!