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 mismatch 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:

First level
The model is y=y0*(1M(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. 
Second level
There is a separate dataset used for this level of the model. The model is now extended to y=y0*(1M(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*(1M), 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**((P6)**2) + n2/10**((P9)**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**((P6)**2) + n2/10**((P9)**2))*xdata/(1+xdata)))
return pm.Normal('y_obs', mu=y0*(1M+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 mismatch. (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:
 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.
 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.
 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.