ArviZ plot_trace does not properly plot multidimensional variables

Question

I'm attempting to run a basic test model using PyMC3, but I've found the ArviZ plot_trace function won't properly show my traces.

Code

from scipy import stats
import arviz as az
import numpy as np
import matplotlib.pyplot as plt
import pymc3 as pm
import seaborn as sns
import pandas as pd
from theano import shared
from sklearn import preprocessing

if __name__ == "__main__":
    with basic_model:

        # Priors for unknown model parameters
        alpha = pm.Normal('alpha', mu=0, sigma=10)
        beta = pm.Normal('beta', mu=0, sigma=10, shape=2)
        sigma = pm.HalfNormal('sigma', sigma=1)

        # Expected value of outcome
        mu = alpha + beta[0]*X1 + beta[1]*X2

        # Likelihood (sampling distribution) of observations
        Y_obs = pm.Normal('Y_obs', mu=mu, sigma=sigma, observed=Y)

        # draw 500 posterior samples
        trace = pm.sample(5000)

    az.plot_trace(trace, compact = False)

The beta parameter is multidimensional, and has both beta[0] and beta[1], but the ArviZ trace only shows beta[0]:

Trace Plot

If I run the trace plot as az.plot_trace(trace, compact = True), then I do see both dimensions of beta properly overlaid. I only observed this issue when trying to plot the dimensions in separate axes with compact = False.

---

Versions

• ArviZ: 0.6.1
• Numpy: 1.18.1
• SciPy: 1.4.1
• xarray: 0.15.0
• Matplotlib: 3.1.3

Answer 1

It looks like you are running into this bug. I'd recommend updating ArviZ to its latest version (0.7.0 at the time of writing) which already includes a fix for this particular bug.

If you had some kind of version constraint, then disabling numba should fix the issue.