Code Consolidation: Same hierarchy on multiple variables

Questions version agnostic
1

Hello Good people of PyMC Discourse,

I have a question for you about code consolidation when you have lots of variables with the same hierarchical structure. I am not sure how to implement the same hierarchical structure across multiple variable in a clean way. For example, let's say I have the following dataset:

import arviz as az
import numpy as np
import pandas as pd
import pymc3 as pm

data = pd.DataFrame({
    "x1":[1, 2, 3, 4, 5]*2,
    "x2":[2, 4, 8, 16, 32]*2,
    "x3":[3, 9, 27, 81, 243]*2,
    "level":[0,0,0,0,0,1,1,1,1,1],
    "y":[14,37,100,279,798,4,9,19,43,107]
})

the variables x1, x2, x3 are going to work as stand ins for much larger dataset that I have. So the goal I have is to build a simple hierarchical regression model predicting on y but have a prior for x1,x2,x3 based on "level".

model_coords = {}
model_coords["obs_id"] = np.arange(len(data))
model_coords["continuous"] = ["x1", "x2", "x3"]
model_coords["level"] = ["0", "1"]

with pm.Model(coords = model_coords) as model:
        
    mu_beta = pm.Normal("mu_beta", 0, 1)
    sigma_beta = pm.HalfNormal("sigma_beta", 1.)
    beta_offset = pm.Normal("beta_offset", 0, 1, dims=("continuous"))
    beta = pm.Deterministic("beta", mu_beta + sigma_beta*beta_offset, dims=("continuous"))
    epsilon = pm.HalfNormal("epsilon", 1)
    
    estimate =  pm.math.dot(data[continuous], beta)
    
    output = pm.Normal("output", mu=estimate, sigma=epsilon, observed = data.y.values, dims = "obs_id" )

using pm.math.dot I’m able to consolidate the summation to just a single function, but how do I do that while also including a hierarchy on beta based on "level"? Also, I realized after playing around with toy samples, that I actually only know how to add a hierarchy to the intercept term. Any input would be greatly appreciated! Thank you in advance!