How to resolve Input Dimension Mis-match Error in Hierarchical Bayesian Inference with PyMC3

Hello everyone,

I am trying to implement a slightly peculiar model using hierarchical bayesian inference with pymc3 and have been encountering a ‘ValueError’ due to input dimension mis-match in Pymc3. The model has two levels of hierarchy with the entire distributions of parameters inferred in the first level being carried over to the second level. The details of the model are as follows:

  1. First level
    The model is y=y0*(1-M(x)). Where M(x) contains two parameters m1 and m2 that are to be inferred. M(x) also contains a constant P which assumes 5 different known constant values. The expression for M(x) is in the code below. y0 is a constant.

  2. Second level
    There is a separate dataset used for this level of the model. The model is now extended to y=y0*(1-M(x)+N(x)). The parameters m1 and m2 used in the above level is utilized here as entire distributions in order to carry over the uncertainty associated with them and is contained in M(x). I am using shared variables for this purpose. There are two new parameters in N(x) called n1 and n2 that are to be inferred in this level. N(x) also contains a constant P which takes 2 different values which are also known. These two new constant values are also used in M(x) at this level.

I have provided here a simplistic version of the code with synthetic data that replicates the error

import numpy as np 
import matplotlib.pyplot as plt
import pymc3 as pm
import theano.tensor as tt
import theano
import arviz as az

np.random.seed(0)

##### Data generation - Level I model    #####
y0 = 100
m1 = 10
m2 = 0.1
xI_list = []
yI_list = []
PI_list = []

for P in [5,7,12,14,15]:
    xI = np.linspace(1,10,10)
    M = np.exp(-m1/P + m2) * xI/(1+xI)
    yI = y0 * (1 - M) + np.random.uniform(low=-1, high=1, size=len(xI))
    yI_list = np.append(yI_list,yI)
    xI_list = np.append(xI_list,xI)
    PI_list = np.append(PI_list, P*np.ones(len(yI)))
#####################################

#### Bayesian inference - Level 1
def likelihoodI(ydata,m1,m2,xdata,Pdata):
    M = tt.exp(-m1/Pdata + m2) * xdata/(1+xdata)
    return pm.Normal('y_obs', mu=y0*(1-M), sigma=1, observed=ydata)

with pm.Model() as first_level_model:
    # Priors for m1, m2
    m1 = pm.Normal('m1', mu=10, sd=1)
    m2 = pm.Normal('m2', mu=0.1, sd=1)

    # Likelihood function call
    likelihoodI(yI_list,m1,m2,xI_list,PI_list)

    # MCMC sampling
    trace1 = pm.sample(4000, tune=2000, cores=2)
        

##### Data generation - Level II model    #####
n1 = 0.4
n2 = 0.6
xII_list = []
yII_list = []
PII_list = []

for P in [6,9]:
    xII = np.linspace(1,10,10)
    M = np.exp(-m1/P + m2) * xII/(1+xII)
    N = (n1/10**((P-6)**2) + n2/10**((P-9)**2))*xII/(1+xII)
    yII = y0 * (1 - M + N) + np.random.uniform(low=-1, high=1, size=len(xI))
    yII_list = np.append(yII_list,yII)
    xII_list = np.append(xII_list,xII)
    PII_list = np.append(PII_list, P*np.ones(len(xII)))
    

#### Bayesian inference - Level 2

# Shared variables for 'm1' and 'm2' to consider their entire distribution in the second level
m1_shared = theano.shared(trace1['m1'])
m2_shared = theano.shared(trace1['m2'])

def likelihoodII(ydata,m1,m2,xdata,Pdata):
    M = tt.exp(-m1/Pdata + m2) * xdata/(1+xdata)
    N = tt.switch(tt.eq(Pdata,6),
                        n1*xdata/(1+xdata), 
                        tt.switch(tt.eq(Pdata,9),
                                  n2*xdata/(1+xdata),
                                  (n1/10**((P-6)**2) + n2/10**((P-9)**2))*xdata/(1+xdata)))
    return pm.Normal('y_obs', mu=y0*(1-M+N), sigma=1, observed=ydata)


with pm.Model() as second_level_model:

    b1 = pm.Normal('b1', mu=0.4, sd=1)
    b2 = pm.Normal('b2', mu=0.6, sd=1)

    # Likelihood function
    likelihoodII(yII_list,m1_shared,m2_shared,xII_list,PII_list)

    # sampling
    trace2 = pm.sample(4000, tune=2000, cores=2)

az.plot_trace(trace2)
plt.show()

When I run the code, I am getting the following error when implementing this in pymc3:

ValueError: Input dimension mis-match. (input[0].shape[0] = 8000, input[1].shape[0] = 20)

How can I resolve this error? Any answers or directions to think would be of great help since I am new to this level of programming. I have also made a similar post in stack overflow since I am stuck on this problem for several weeks now.

Some pointers:

  1. I understand that the error is caused due to the size mismatch between the input data (xdata/ydata) for level 2 and m1_shared and m2_shared coming from the posterior sampling of level 1. The size of the parameters are depending on the sampling size while the size of the input data depends on the dataset. So I do not have much control over the sizes of these.
  2. Due to the different values taken by P, I need to use some form of conditional statements in the likelihood formulation. I am using theano here which is partly causing the problem in my opinion.
  3. Since theano is not supported anymore, I would also be happy to migrate to jax if the same can be done with jax. If someone could point me in that direction, it would also be very kind of them.

First, I suggest you use the latest PyMC (not the quite old PyMC3). It uses pytensor instead of theano, which is still actively maintained by the PyMC developers.

Second, I suggest you print variable.eval().shape in a couple of places to see if the shapes you are getting are what you expect. Theano/PyTensor uses the same rules as Numpy when it comes to shape handling and broadcasting.

1 Like

This may also help you understanding shape issues that happen at the PyMC distribution level: Distribution Dimensionality — PyMC 5.10.2 documentation

1 Like

Thanks a lot. I’ll try to make these changes and get back here. I do have one more question which I have posted as a separate query. Would be very nice if you guys could take a look at it. Cheers!