Complicated, long polynomial expressions in Model - precompiling?

In my model I use some complicated polynomial functions - actually multivariate polynomial P(v,w)(parameters v and w). It takes quite a while to compile functions when I start the script (before sampling). [I shortened expressions substantially in example code bellow]

Is there any way to speed up this pre-sampling step by somehow pre-compiling such functions? What is the best way to approach inference of such models?

with pm.Model() as model:    
    # Prior distributions
    dg = pm.Normal('dg', mu=-0.20, sigma=0.2)
    dh = pm.Normal('dh', mu=-1.0, sigma=0.2)
    dcp = pm.Normal('dcp', mu=0, sigma=0.01)
    v = pm.Normal('v', mu=0.06, sigma=0.02)
    w = np.exp((-1 / R) * (dg / T0 + dh * ((1 / T_CD_K) - (1 / T0)) + dcp * (1 - (T0 / T_CD_K) - pm.math.log(T_CD_K/ T0))))
    
    # More priors
    H2=pm.Normal('H2', mu=-40000, sigma=5000)
    dH2t=pm.Normal('dH2', mu=1000, sigma=300)
    C=pm.Normal('C', mu=2200, sigma=500)
    C_temp=pm.Normal('C_temp', mu=-50, sigma=50)
    k=pm.Normal('k', mu=3.8, sigma=1)
    
    # Deterministic variables
    H2_t=H2 + dH2t*T_CD
    H1_t=H2_t*(1-k/6)
    C_t=C+C_temp*T_CD
    
    # Get model functions from another script
    Q=LR_functions(w,v,H1_t,H2_t,C_t)["Q_"+str(n_sim)]
    cd=LR_functions(w,v,H1_t,H2_t,C_t)["cd_"+str(n_sim)]
    
    #Explicit expression of model fucntion - polynomial of v,w
    #Q=54*v**6*w**8 + 240*v**6*w**7 + 630*v**6*w**6 + 1260*v**6*w**5 + 2100*v**6*w**4 + 3024*v**6*w**3 + 3780*v**6*w**2 + 3960*v**6*w + 33*v**5*w**10 + 120*v**5*w**9 + 270*v**5*w**8 + 480*v**5*w**7 + 735*v**5*w**6 + 1008*v**5*w**5 + 1260*v**5*w**4 + 1440*v**5*w**3 + 1485*v**5*w**2 + 1320*v**5*w + 13*v**4*w**12 + 39*v**4*w**11 + 78*v**4*w**10 + 130*v**4*w**9 + 195*v**4*w**8 + 273*v**4*w**7 + ...
    #cd=19656*C*v**4 + 3780*C*v**3 + 2448*C*v**2 + 306*C*v + 18*C + v**6*w**8*(84*C + 72*H1 + 60*H2) + v**6*w**8*(168.0*C + 472.5*H1 + 115.5*H2) + v**6*w**7*(480*C + 360*H1 + 240*H2) + v**6*w**7*(900.0*C + 1980.0*H1 + 360.0*H2) + v**6*w**6*(1620*C + 1080*H1 + 540*H2) + v**6*w**6*(2700.0*C + 4950.0*H1 + 450.0*H2) + v**6*w**5*(4200*C + 2520*H1 + 840*H2) + v**6*w**5*(5880.0*C + 8820.0*H1 + 420.0*H2) + ...
    
    mu = cd/(Q*(n_sim+1)) #final model calc
    
    # Likelihood
    sd = pm.HalfNormal('sd', sigma=100)
    obs = pm.Normal('obs', mu=mu, sigma=sd, observed=Y)

    # Perform MCMC
    trace = pm.sample(3000, tune=1000,nuts_sampler= "numpyro")

You can perhaps have a look at what those forms get rewritten into and write it already like that.

Does the compilation still take a lot of time in subsequent calls? The C code gets cached after the first time so it should be faster if the bottleneck is the compilation and not the graph rewriting itself.

If I print function (pytensor.dprint) I get this pretty complicated computational graph bellow. This graph is generated extremely fast even for some really long, complicated expressions. Can I use such graph in any way?

Actually I don’t notice any speed up (similiar timings), when calling script again. (I’m calling scripts on Ubuntu20.04 via WSL )

Add [id A]
 β”œβ”€ Add [id B]
 β”‚  β”œβ”€ Add [id C]
 β”‚  β”‚  β”œβ”€ Add [id D]
 β”‚  β”‚  β”‚  β”œβ”€ Add [id E]
 β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id F]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id G]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id H]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id I]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id J]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id K]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id L]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id M]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Add [id N]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ Mul [id O]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”œβ”€ ExpandDims{axis=0} [id P]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id Q]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ 3 [id R]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Pow [id S]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  β”œβ”€ RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F7ABB8FE5E0>) [id U]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  β”œβ”€ [] [id V]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  β”œβ”€ 11 [id W]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  β”œβ”€ 0.06 [id X]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  └─ 0.02 [id Y]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ 4 [id Z]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Pow [id BA]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id BC]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ ExpandDims{axis=0} [id BD]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”‚  └─ -503.27126321087064 [id BE]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Add [id BF]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ Add [id BG]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”œβ”€ ExpandDims{axis=0} [id BH]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚  └─ True_div [id BI]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id BJ] 'dg'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     β”‚  β”œβ”€ RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F7ABB8FE420>) [id BK]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     β”‚  β”œβ”€ [] [id BL]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     β”‚  β”œβ”€ 11 [id BM]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     β”‚  β”œβ”€ -0.2 [id BN]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     β”‚  └─ 0.2 [id BO]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  β”‚     └─ 273.15 [id BP]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Mul [id BQ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”œβ”€ ExpandDims{axis=0} [id BR]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”‚  └─ normal_rv{0, (0, 0), floatX, False}.1 [id BS] 'dh'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”‚     β”œβ”€ RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F7ABB8FEDC0>) [id BT]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”‚     β”œβ”€ [] [id BU]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”‚     β”œβ”€ 11 [id BV]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”‚     β”œβ”€ -1.0 [id BW]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     β”‚     └─ 0.2 [id BX]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚     └─ [ 0.000000 ... 04684e-04] [id BY]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ Mul [id BZ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”œβ”€ ExpandDims{axis=0} [id CA]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”‚  └─ normal_rv{0, (0, 0), floatX, False}.1 [id CB] 'dcp'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”‚     β”œβ”€ RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7F7ABB8FF680>) [id CC]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”‚     β”œβ”€ [] [id CD]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”‚     β”œβ”€ 11 [id CE]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”‚     β”œβ”€ 0 [id CF]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           β”‚     └─ 0.01 [id CG]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚           └─ Sub [id CH]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚              β”œβ”€ [0. ... .26798874] [id CI]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚              └─ Log [id CJ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚                 └─ [1. ... .36609921] [id CK]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ ExpandDims{axis=0} [id CL]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ 2 [id CM]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id CN]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id CO]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id CP]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 9 [id CQ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id CR]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ 4 [id CS]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ ExpandDims{axis=0} [id CT]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Mul [id CU]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ 12 [id CV]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ Pow [id CW]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           └─ 4 [id CX]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id CY]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id CZ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id DA]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 2 [id DB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id DC]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ 3 [id DD]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Pow [id DE]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ ExpandDims{axis=0} [id DF]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           └─ 3 [id DG]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id DH]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id DI]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id DJ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 6 [id DK]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id DL]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ 3 [id DM]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Pow [id DN]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ ExpandDims{axis=0} [id DO]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           └─ 2 [id DP]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id DQ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id DR]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id DS]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 12 [id DT]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id DU]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ 3 [id DV]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ ExpandDims{axis=0} [id DW]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Mul [id DX]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ 20 [id DY]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ Pow [id DZ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           └─ 3 [id EA]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id EB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id EC]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Pow [id ED]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ 2 [id EE]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Pow [id EF]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ ExpandDims{axis=0} [id EG]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           └─ 5 [id EH]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id EI]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id EJ]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id EK]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 2 [id EL]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id EM]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ 2 [id EN]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚     └─ Pow [id EO]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚        └─ ExpandDims{axis=0} [id EP]
 β”‚  β”‚  β”‚  β”‚  β”‚  β”‚           └─ 4 [id EQ]
 β”‚  β”‚  β”‚  β”‚  β”‚  └─ Mul [id ER]
 β”‚  β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id ES]
 β”‚  β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id ET]
 β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 3 [id EU]
 β”‚  β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id EV]
 β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚     β”‚        └─ 2 [id EW]
 β”‚  β”‚  β”‚  β”‚  β”‚     └─ Pow [id EX]
 β”‚  β”‚  β”‚  β”‚  β”‚        β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚  β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚  β”‚        └─ ExpandDims{axis=0} [id EY]
 β”‚  β”‚  β”‚  β”‚  β”‚           └─ 3 [id EZ]
 β”‚  β”‚  β”‚  β”‚  └─ Mul [id FA]
 β”‚  β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id FB]
 β”‚  β”‚  β”‚  β”‚     β”‚  └─ Mul [id FC]
 β”‚  β”‚  β”‚  β”‚     β”‚     β”œβ”€ 4 [id FD]
 β”‚  β”‚  β”‚  β”‚     β”‚     └─ Pow [id FE]
 β”‚  β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚     β”‚        └─ 2 [id FF]
 β”‚  β”‚  β”‚  β”‚     └─ Pow [id FG]
 β”‚  β”‚  β”‚  β”‚        β”œβ”€ Exp [id BB]
 β”‚  β”‚  β”‚  β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚  β”‚        └─ ExpandDims{axis=0} [id FH]
 β”‚  β”‚  β”‚  β”‚           └─ 2 [id FI]
 β”‚  β”‚  β”‚  └─ Mul [id FJ]
 β”‚  β”‚  β”‚     β”œβ”€ ExpandDims{axis=0} [id FK]
 β”‚  β”‚  β”‚     β”‚  └─ Mul [id FL]
 β”‚  β”‚  β”‚     β”‚     β”œβ”€ 5 [id FM]
 β”‚  β”‚  β”‚     β”‚     └─ Pow [id FN]
 β”‚  β”‚  β”‚     β”‚        β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚  β”‚     β”‚        β”‚  └─ Β·Β·Β·
 β”‚  β”‚  β”‚     β”‚        └─ 2 [id FO]
 β”‚  β”‚  β”‚     └─ Exp [id BB]
 β”‚  β”‚  β”‚        └─ Β·Β·Β·
 β”‚  β”‚  └─ ExpandDims{axis=0} [id FP]
 β”‚  β”‚     └─ Mul [id FQ]
 β”‚  β”‚        β”œβ”€ 21 [id FR]
 β”‚  β”‚        └─ Pow [id FS]
 β”‚  β”‚           β”œβ”€ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚  β”‚           β”‚  └─ Β·Β·Β·
 β”‚  β”‚           └─ 2 [id FT]
 β”‚  └─ ExpandDims{axis=0} [id FU]
 β”‚     └─ Mul [id FV]
 β”‚        β”œβ”€ 7 [id FW]
 β”‚        └─ normal_rv{0, (0, 0), floatX, False}.1 [id T] 'v'
 β”‚           └─ Β·Β·Β·
 └─ ExpandDims{axis=0} [id FX]
    └─ 1 [id FY]

You want to print the graph of model.compile_logp().f instead

This is the result I obtain when I print β€˜model.compile_logp().f’.

Sum{axes=None} [id A] '__logp' 31
 └─ MakeVector{dtype='float64'} [id B] 30
    β”œβ”€ Composite{(i4 + (i3 * sqr((i2 * (i0 + i1)))))} [id C] 'sigma > 0' 29
    β”‚  β”œβ”€ 0.2 [id D]
    β”‚  β”œβ”€ dg [id E]
    β”‚  β”œβ”€ 5.0 [id F]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 0.6904993792294276 [id H]
    β”œβ”€ Composite{(i4 + (i3 * sqr((i2 * (i0 + i1)))))} [id I] 'sigma > 0' 28
    β”‚  β”œβ”€ 1.0 [id J]
    β”‚  β”œβ”€ dh [id K]
    β”‚  β”œβ”€ 5.0 [id F]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 0.6904993792294276 [id H]
    β”œβ”€ Composite{(i3 + (i2 * sqr((i0 * i1))))} [id L] 'sigma > 0' 27
    β”‚  β”œβ”€ 100.0 [id M]
    β”‚  β”œβ”€ dcp [id N]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 3.6862316527834182 [id O]
    β”œβ”€ Composite{(i4 + (i3 * sqr((i2 * (i0 - i1)))))} [id P] 'sigma > 0' 26
    β”‚  β”œβ”€ v [id Q]
    β”‚  β”œβ”€ 0.06 [id R]
    β”‚  β”œβ”€ 50.0 [id S]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 2.9930844722234733 [id T]
    β”œβ”€ Composite{((i3 * sqr((i2 * (i0 + i1)))) - i4)} [id U] 'sigma > 0' 25
    β”‚  β”œβ”€ 40000.0 [id V]
    β”‚  β”œβ”€ H2 [id W]
    β”‚  β”œβ”€ 0.0002 [id X]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 9.43613172462091 [id Y]
    β”œβ”€ Composite{((i3 * sqr((i2 * (i0 - i1)))) - i4)} [id Z] 'sigma > 0' 24
    β”‚  β”œβ”€ dH2 [id BA]
    β”‚  β”œβ”€ 1000.0 [id BB]
    β”‚  β”œβ”€ 0.00333333 ... 3333333335 [id BC]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 6.622721007860873 [id BD]
    β”œβ”€ Composite{((i3 * sqr((i2 * (i0 - i1)))) - i4)} [id BE] 'sigma > 0' 23
    β”‚  β”œβ”€ C [id BF]
    β”‚  β”œβ”€ 2200.0 [id BG]
    β”‚  β”œβ”€ 0.002 [id BH]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 7.133546631626864 [id BI]
    β”œβ”€ Composite{((i3 * sqr((i2 * (i0 + i1)))) - i4)} [id BJ] 'sigma > 0' 22
    β”‚  β”œβ”€ 50.0 [id S]
    β”‚  β”œβ”€ C_temp [id BK]
    β”‚  β”œβ”€ 0.02 [id BL]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 4.830961538632819 [id BM]
    β”œβ”€ Composite{((i2 * sqr((i0 - i1))) - i3)} [id BN] 'sigma > 0' 21
    β”‚  β”œβ”€ k [id BO]
    β”‚  β”œβ”€ 3.8 [id BP]
    β”‚  β”œβ”€ -0.5 [id G]
    β”‚  └─ 0.9189385332046727 [id BQ]
    β”œβ”€ Composite{...}.1 [id BR] 'sd_log___logprob' 0
    β”‚  β”œβ”€ sd_log__ [id BS]
    β”‚  β”œβ”€ 0.00333333 ... 3333333335 [id BC]
    β”‚  β”œβ”€ 0.0 [id BT]
    β”‚  β”œβ”€ -inf [id BU]
    β”‚  β”œβ”€ 5.929573827300929 [id BV]
    β”‚  └─ -0.5 [id G]
    └─ Sum{axes=None} [id BW] 20
       └─ Composite{...} [id BX] 'sigma > 0' 19
          β”œβ”€ ExpandDims{axis=0} [id BY] 18
          β”‚  └─ dcp [id N]
          β”œβ”€ [ 0.000000 ... 06443e-02] [id BZ]
          β”œβ”€ Mul [id CA] 17
          β”‚  β”œβ”€ [0.00366099] [id CB]
          β”‚  └─ ExpandDims{axis=0} [id CC] 16
          β”‚     └─ dg [id E]
          β”œβ”€ ExpandDims{axis=0} [id CD] 15
          β”‚  └─ dh [id K]
          β”œβ”€ [ 0.000000 ... 04684e-04] [id CE]
          β”œβ”€ [-503.27126321] [id CF]
          β”œβ”€ [5.] [id CG]
          β”œβ”€ Composite{...}.0 [id CH] 8
          β”‚  β”œβ”€ ExpandDims{axis=0} [id CI] 4
          β”‚  β”‚  └─ v [id Q]
          β”‚  └─ [21.] [id CJ]
          β”œβ”€ [12.] [id CK]
          β”œβ”€ Composite{...}.0 [id CL] 9
          β”‚  β”œβ”€ ExpandDims{axis=0} [id CI] 4
          β”‚  β”‚  └─ Β·Β·Β·
          β”‚  └─ [20.] [id CM]
          β”œβ”€ [9.] [id CN]
          β”œβ”€ Composite{...}.0 [id CO] 6
          β”‚  β”œβ”€ ExpandDims{axis=0} [id CI] 4
          β”‚  β”‚  └─ Β·Β·Β·
          β”‚  β”œβ”€ [0.16666667] [id CP]
          β”‚  β”œβ”€ ExpandDims{axis=0} [id CQ] 5
          β”‚  β”‚  └─ k [id BO]
          β”‚  └─ [12.] [id CK]
          β”œβ”€ [20.] [id CM]
          β”œβ”€ ExpandDims{axis=0} [id CR] 14
          β”‚  └─ C [id BF]
          β”œβ”€ Composite{(i2 - (i0 * i1))} [id CS] 13
          β”‚  β”œβ”€ [0.16666667] [id CP]
          β”‚  β”œβ”€ ExpandDims{axis=0} [id CQ] 5
          β”‚  β”‚  └─ Β·Β·Β·
          β”‚  └─ [1.] [id CT]
          β”œβ”€ ExpandDims{axis=0} [id CU] 12
          β”‚  └─ H2 [id W]
          β”œβ”€ ExpandDims{axis=0} [id CV] 11
          β”‚  └─ dH2 [id BA]
          β”œβ”€ [  0. ... .        ] [id CW]
          β”œβ”€ ExpandDims{axis=0} [id CX] 10
          β”‚  └─ C_temp [id BK]
          β”œβ”€ [48.] [id CY]
          β”œβ”€ [36.] [id CZ]
          β”œβ”€ [1.] [id CT]
          β”œβ”€ Composite{...}.1 [id CO] 6
          β”‚  └─ Β·Β·Β·
          β”œβ”€ Composite{...}.1 [id CL] 9
          β”‚  └─ Β·Β·Β·
          β”œβ”€ Composite{...}.1 [id CH] 8
          β”‚  └─ Β·Β·Β·
          β”œβ”€ Mul [id DA] 7
          β”‚  β”œβ”€ [7.] [id DB]
          β”‚  └─ ExpandDims{axis=0} [id CI] 4
          β”‚     └─ Β·Β·Β·
          β”œβ”€ [4.] [id DC]
          β”œβ”€ [3.] [id DD]
          β”œβ”€ [2.] [id DE]
          β”œβ”€ [6.] [id DF]
          β”œβ”€ [8.] [id DG]
          β”œβ”€ Composite{...}.2 [id CO] 6
          β”‚  └─ Β·Β·Β·
          β”œβ”€ [56.] [id DH]
          β”œβ”€ ExpandDims{axis=0} [id CI] 4
          β”‚  └─ Β·Β·Β·
          β”œβ”€ [168.] [id DI]
          β”œβ”€ [160.] [id DJ]
          β”œβ”€ [96.] [id DK]
          β”œβ”€ [0.125] [id DL]
          β”œβ”€ obs{[1147.3534 ... .31715505]} [id DM]
          β”œβ”€ ExpandDims{axis=0} [id DN] 1
          β”‚  └─ Composite{...}.0 [id BR] 0
          β”‚     └─ Β·Β·Β·
          β”œβ”€ [-0.5] [id DO]
          β”œβ”€ [0.91893853] [id DP]
          β”œβ”€ Log [id DQ] 3
          β”‚  └─ ExpandDims{axis=0} [id DN] 1
          β”‚     └─ Β·Β·Β·
          β”œβ”€ Gt [id DR] 2
          β”‚  β”œβ”€ ExpandDims{axis=0} [id DN] 1
          β”‚  β”‚  └─ Β·Β·Β·
          β”‚  └─ [0] [id DS]
          └─ [-inf] [id DT]

Inner graphs:

Composite{(i4 + (i3 * sqr((i2 * (i0 + i1)))))} [id C]
 ← add [id DU] 'o0'
    β”œβ”€ i4 [id DV]
    └─ mul [id DW]
       β”œβ”€ i3 [id DX]
       └─ sqr [id DY]
          └─ mul [id DZ]
             β”œβ”€ i2 [id EA]
             └─ add [id EB]
                β”œβ”€ i0 [id EC]
                └─ i1 [id ED]

Composite{(i4 + (i3 * sqr((i2 * (i0 + i1)))))} [id I]
 ← add [id EE] 'o0'
    β”œβ”€ i4 [id EF]
    └─ mul [id EG]
       β”œβ”€ i3 [id EH]
       └─ sqr [id EI]
          └─ mul [id EJ]
             β”œβ”€ i2 [id EK]
             └─ add [id EL]
                β”œβ”€ i0 [id EM]
                └─ i1 [id EN]

Composite{(i3 + (i2 * sqr((i0 * i1))))} [id L]
 ← add [id EO] 'o0'
    β”œβ”€ i3 [id EP]
    └─ mul [id EQ]
       β”œβ”€ i2 [id ER]
       └─ sqr [id ES]
          └─ mul [id ET]
             β”œβ”€ i0 [id EU]
             └─ i1 [id EV]

Composite{(i4 + (i3 * sqr((i2 * (i0 - i1)))))} [id P]
 ← add [id EW] 'o0'
    β”œβ”€ i4 [id EX]
    └─ mul [id EY]
       β”œβ”€ i3 [id EZ]
       └─ sqr [id FA]
          └─ mul [id FB]
             β”œβ”€ i2 [id FC]
             └─ sub [id FD]
                β”œβ”€ i0 [id FE]
                └─ i1 [id FF]

Composite{((i3 * sqr((i2 * (i0 + i1)))) - i4)} [id U]
 ← sub [id FG] 'o0'
    β”œβ”€ mul [id FH]
    β”‚  β”œβ”€ i3 [id FI]
    β”‚  └─ sqr [id FJ]
    β”‚     └─ mul [id FK]
    β”‚        β”œβ”€ i2 [id FL]
    β”‚        └─ add [id FM]
    β”‚           β”œβ”€ i0 [id FN]
    β”‚           └─ i1 [id FO]
    └─ i4 [id FP]

Composite{((i3 * sqr((i2 * (i0 - i1)))) - i4)} [id Z]
 ← sub [id FQ] 'o0'
    β”œβ”€ mul [id FR]
    β”‚  β”œβ”€ i3 [id FS]
    β”‚  └─ sqr [id FT]
    β”‚     └─ mul [id FU]
    β”‚        β”œβ”€ i2 [id FV]
    β”‚        └─ sub [id FW]
    β”‚           β”œβ”€ i0 [id FX]
    β”‚           └─ i1 [id FY]
    └─ i4 [id FZ]

Composite{((i3 * sqr((i2 * (i0 - i1)))) - i4)} [id BE]
 ← sub [id GA] 'o0'
    β”œβ”€ mul [id GB]
    β”‚  β”œβ”€ i3 [id GC]
    β”‚  └─ sqr [id GD]
    β”‚     └─ mul [id GE]
    β”‚        β”œβ”€ i2 [id GF]
    β”‚        └─ sub [id GG]
    β”‚           β”œβ”€ i0 [id GH]
    β”‚           └─ i1 [id GI]
    └─ i4 [id GJ]

Composite{((i3 * sqr((i2 * (i0 + i1)))) - i4)} [id BJ]
 ← sub [id GK] 'o0'
    β”œβ”€ mul [id GL]
    β”‚  β”œβ”€ i3 [id GM]
    β”‚  └─ sqr [id GN]
    β”‚     └─ mul [id GO]
    β”‚        β”œβ”€ i2 [id GP]
    β”‚        └─ add [id GQ]
    β”‚           β”œβ”€ i0 [id GR]
    β”‚           └─ i1 [id GS]
    └─ i4 [id GT]

Composite{((i2 * sqr((i0 - i1))) - i3)} [id BN]
 ← sub [id GU] 'o0'
    β”œβ”€ mul [id GV]
    β”‚  β”œβ”€ i2 [id GW]
    β”‚  └─ sqr [id GX]
    β”‚     └─ sub [id GY]
    β”‚        β”œβ”€ i0 [id GZ]
    β”‚        └─ i1 [id HA]
    └─ i3 [id HB]

Composite{...} [id BR]
 ← exp [id HC] 'o0'
    └─ i0 [id HD]
 ← add [id HE] 'o1'
    β”œβ”€ Switch [id HF]
    β”‚  β”œβ”€ GE [id HG]
    β”‚  β”‚  β”œβ”€ exp [id HC] 'o0'
    β”‚  β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  └─ i2 [id HH]
    β”‚  β”œβ”€ sub [id HI]
    β”‚  β”‚  β”œβ”€ mul [id HJ]
    β”‚  β”‚  β”‚  β”œβ”€ i5 [id HK]
    β”‚  β”‚  β”‚  └─ sqr [id HL]
    β”‚  β”‚  β”‚     └─ mul [id HM]
    β”‚  β”‚  β”‚        β”œβ”€ i1 [id HN]
    β”‚  β”‚  β”‚        └─ exp [id HC] 'o0'
    β”‚  β”‚  β”‚           └─ Β·Β·Β·
    β”‚  β”‚  └─ i4 [id HO]
    β”‚  └─ i3 [id HP]
    └─ i0 [id HD]

Composite{...} [id BX]
 ← Switch [id HQ] 'o0'
    β”œβ”€ i43 [id HR]
    β”œβ”€ sub [id HS]
    β”‚  β”œβ”€ sub [id HT]
    β”‚  β”‚  β”œβ”€ mul [id HU]
    β”‚  β”‚  β”‚  β”œβ”€ i40 [id HV]
    β”‚  β”‚  β”‚  └─ sqr [id HW]
    β”‚  β”‚  β”‚     └─ true_div [id HX]
    β”‚  β”‚  β”‚        β”œβ”€ sub [id HY]
    β”‚  β”‚  β”‚        β”‚  β”œβ”€ i38 [id HZ]
    β”‚  β”‚  β”‚        β”‚  └─ true_div [id IA]
    β”‚  β”‚  β”‚        β”‚     β”œβ”€ mul [id IB]
    β”‚  β”‚  β”‚        β”‚     β”‚  β”œβ”€ i37 [id IC]
    β”‚  β”‚  β”‚        β”‚     β”‚  └─ add [id ID]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id IE]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i36 [id IF]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ add [id IG] 't74'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  β”œβ”€ i13 [id IH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ mul [id II] 't11'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     β”œβ”€ i18 [id IJ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     └─ i17 [id IK]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ i11 [id IL]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id IM]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i35 [id IN]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ add [id IG] 't74'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id IP]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i34 [id IQ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ add [id IG] 't74'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id IS]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i32 [id IT]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ add [id IG] 't74'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ i33 [id IU]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id IV]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i30 [id IW]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id IG] 't74'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id IX]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i30 [id IW]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ add [id IY] 't35'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  β”œβ”€ i15 [id IZ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ mul [id JA] 't44'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     β”œβ”€ i16 [id JB]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     └─ i17 [id IK]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i31 [id JC]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        └─ mul [id JF]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚           β”œβ”€ i5 [id JG]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚           └─ add [id JH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚              β”œβ”€ i2 [id JI]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚              β”œβ”€ mul [id JJ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚              β”‚  β”œβ”€ i3 [id JK]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚              β”‚  └─ i4 [id JL]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚              └─ mul [id JM]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚                 β”œβ”€ i0 [id JN]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚                 └─ i1 [id JO]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id JP]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i28 [id JQ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i11 [id IL]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id JR] 't49'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ mul [id JS]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  β”œβ”€ i27 [id JT]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  └─ add [id IG] 't74'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ mul [id JU]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”œβ”€ i6 [id JV]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”œβ”€ add [id IY] 't35'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        └─ i14 [id JW]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id JX]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i20 [id JY]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i11 [id IL]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id JZ] 't104'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ i13 [id IH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ mul [id II] 't11'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ mul [id KA]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”œβ”€ add [id IY] 't35'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        └─ i14 [id JW]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KB]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i26 [id KC]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ Composite{(sqr(i0) * i0)} [id KD] 't78'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id KE] 't14'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ i13 [id IH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ mul [id II] 't11'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ mul [id KF] 't27'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”œβ”€ i27 [id JT]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”œβ”€ add [id IY] 't35'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        └─ i14 [id JW]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KG]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i29 [id KH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id JR] 't49'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KI]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i19 [id KJ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id JZ] 't104'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KK]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i28 [id JQ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ Composite{(sqr(sqr(i0)) * i0)} [id KL] 't53'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id KM]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ mul [id KF] 't27'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ i15 [id IZ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ mul [id JA] 't44'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KN]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i28 [id JQ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ Composite{sqr(sqr(i0))} [id KO] 't15'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id KP]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ i13 [id IH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ mul [id II] 't11'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ mul [id KQ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  β”œβ”€ i29 [id KH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  β”œβ”€ add [id IY] 't35'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”‚  └─ i14 [id JW]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     β”œβ”€ i15 [id IZ]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ mul [id JA] 't44'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚        └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KR]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i29 [id KH]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ Composite{(sqr(i0) * i0)} [id KD] 't78'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id KE] 't14'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”œβ”€ mul [id KS]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i26 [id KC]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”œβ”€ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚  └─ add [id JR] 't49'
    β”‚  β”‚  β”‚        β”‚     β”‚     β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚     └─ mul [id KT]
    β”‚  β”‚  β”‚        β”‚     β”‚        β”œβ”€ i12 [id KU]
    β”‚  β”‚  β”‚        β”‚     β”‚        β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚     β”‚        β”œβ”€ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚     β”‚        β”‚  └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     β”‚        └─ add [id JZ] 't104'
    β”‚  β”‚  β”‚        β”‚     β”‚           └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚     └─ add [id KV]
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ i21 [id KW]
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id KX]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i27 [id JT]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i11 [id IL]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id KY]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i10 [id KZ]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i11 [id IL]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ i22 [id LA]
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LB]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i28 [id JQ]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ Composite{(sqr(i0) * i0)} [id KD] 't78'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LC]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i29 [id KH]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LD]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i8 [id LE]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i9 [id IO]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ i23 [id LF]
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LG]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ Composite{(sqr(sqr(i0)) * i0)} [id KL] 't53'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LH]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i28 [id JQ]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ Composite{sqr(sqr(i0))} [id KO] 't15'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LI]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i27 [id JT]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ Composite{(sqr(i0) * i0)} [id KD] 't78'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LJ]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i26 [id KC]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ sqr [id JD] 't24'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ mul [id LK]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i6 [id JV]
    β”‚  β”‚  β”‚        β”‚        β”‚  β”œβ”€ i7 [id IR]
    β”‚  β”‚  β”‚        β”‚        β”‚  └─ exp [id JE] 't87'
    β”‚  β”‚  β”‚        β”‚        β”‚     └─ Β·Β·Β·
    β”‚  β”‚  β”‚        β”‚        β”œβ”€ i24 [id LL]
    β”‚  β”‚  β”‚        β”‚        └─ i25 [id LM]
    β”‚  β”‚  β”‚        └─ i39 [id LN]
    β”‚  β”‚  └─ i41 [id LO]
    β”‚  └─ i42 [id LP]
    └─ i44 [id LQ]

Composite{...} [id CH]
 ← sqr [id LR] 'o0'
    └─ i0 [id LS]
 ← mul [id LT] 'o1'
    β”œβ”€ i1 [id LU]
    └─ sqr [id LR] 'o0'
       └─ Β·Β·Β·

Composite{...} [id CL]
 ← Composite{(sqr(i0) * i0)} [id LV] 'o0'
    └─ i0 [id LW]
 ← mul [id LX] 'o1'
    β”œβ”€ i1 [id LY]
    └─ Composite{(sqr(i0) * i0)} [id LV] 'o0'
       └─ Β·Β·Β·

Composite{...} [id CO]
 ← Composite{sqr(sqr(i0))} [id LZ] 'o0'
    └─ i0 [id MA]
 ← mul [id MB] 'o1'
    β”œβ”€ i3 [id MC]
    └─ Composite{sqr(sqr(i0))} [id LZ] 'o0'
       └─ Β·Β·Β·
 ← sub [id MD] 'o2'
    β”œβ”€ Composite{sqr(sqr(i0))} [id LZ] 'o0'
    β”‚  └─ Β·Β·Β·
    └─ mul [id ME]
       β”œβ”€ i1 [id MF]
       β”œβ”€ i2 [id MG]
       └─ Composite{sqr(sqr(i0))} [id LZ] 'o0'
          └─ Β·Β·Β·

Composite{(i2 - (i0 * i1))} [id CS]
 ← sub [id MH] 'o0'
    β”œβ”€ i2 [id MI]
    └─ mul [id MJ]
       β”œβ”€ i0 [id MK]
       └─ i1 [id ML]

Composite{(sqr(i0) * i0)} [id KD]
 ← mul [id MM] 'o0'
    β”œβ”€ sqr [id MN]
    β”‚  └─ i0 [id MO]
    └─ i0 [id MO]

Composite{(sqr(sqr(i0)) * i0)} [id KL]
 ← mul [id MP] 'o0'
    β”œβ”€ sqr [id MQ]
    β”‚  └─ sqr [id MR]
    β”‚     └─ i0 [id MS]
    └─ i0 [id MS]

Composite{sqr(sqr(i0))} [id KO]
 ← sqr [id MT] 'o0'
    └─ sqr [id MU]
       └─ i0 [id MV]

Composite{(sqr(i0) * i0)} [id LV]
 ← mul [id MW] 'o0'
    β”œβ”€ sqr [id MX]
    β”‚  └─ i0 [id MY]
    └─ i0 [id MY]

Composite{sqr(sqr(i0))} [id LZ]
 ← sqr [id MZ] 'o0'
    └─ sqr [id NA]
       └─ i0 [id NB]

Does it take a long time for PyMC to return compile_logp after the first time? If not, I wouldn’t say there’s an issue with compilation time.

Your model doesn’t seem that bad

I agree with you maybe it is really not compiling. For one quite simple poynomial it takes around 1sec to return compile_logp (that timing persist also in subsequent runs). Also for some more complicated, longer polynomial I am still around 5sec.

My script spends most of the time when executing β€œpm.sample” line, but before entering compilation. Compilation time according to print of NUTS sampler is around 20sec (for that 5sec compile_logp example). I think that’s quite reasonable.

Right before compilation (after a long period of waiting time) I also get this warning…

/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/tensor/rewriting/elemwise.py:1019: UserWarning: Loop fusion failed because the resulting node would exceed the kernel argument limit.
  warn(

Additionally I ran to a compilation error…

You can find the C code in this temporary file: /tmp/pytensor_compilation_error_j21u7nk0
Traceback (most recent call last):
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/vm.py", line 1243, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/op.py", line 131, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/op.py", line 96, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1200, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1120, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1644, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1240, in module_from_key
    module = lnk.compile_cmodule(location)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1543, in compile_cmodule
    module = c_compiler.compile_str(
             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 2649, in compile_str
    raise CompileError(
pytensor.link.c.exceptions.CompileError: Compilation failed (return status=1):
/home/ubufux/anaconda3/envs/pymc_env/bin/g++ -shared -g -O3 -fno-math-errno -Wno-unused-label -Wno-unused-variable -Wno-write-strings -Wno-c++11-narrowing -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables -march=skylake -mmmx -mpopcnt -msse -msse2 -msse3 -mssse3 -msse4.1 -msse4.2 -mavx -mavx2 -mno-sse4a -mno-fma4 -mno-xop -mfma -mno-avx512f -mbmi -mbmi2 -maes -mpclmul -mno-avx512vl -mno-avx512bw -mno-avx512dq -mno-avx512cd -mno-avx512er -mno-avx512pf -mno-avx512vbmi -mno-avx512ifma -mno-avx5124vnniw -mno-avx5124fmaps -mno-avx512vpopcntdq -mno-avx512vbmi2 -mno-gfni -mno-vpclmulqdq -mno-avx512vnni -mno-avx512bitalg -mno-avx512bf16 -mno-avx512vp2intersect -mno-3dnow -madx -mabm -mno-cldemote -mclflushopt -mno-clwb -mno-clzero -mcx16 -mno-enqcmd -mf16c -mfsgsbase -mfxsr -mno-hle -msahf -mno-lwp -mlzcnt -mmovbe -mno-movdir64b -mno-movdiri -mno-mwaitx -mno-pconfig -mno-pku -mno-prefetchwt1 -mprfchw -mno-ptwrite -mno-rdpid -mrdrnd -mrdseed -mno-rtm -mno-serialize -mno-sgx -mno-sha -mno-shstk -mno-tbm -mno-tsxldtrk -mno-vaes -mno-waitpkg -mno-wbnoinvd -mxsave -mxsavec -mxsaveopt -mxsaves -mno-amx-tile -mno-amx-int8 -mno-amx-bf16 -mno-uintr -mno-hreset -mno-kl -mno-widekl -mno-avxvnni -mno-avx512fp16 --param l1-cache-size=32 --param l1-cache-line-size=64 --param l2-cache-size=8192 -mtune=skylake -DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION -m64 -fPIC -I/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/numpy/core/include -I/home/ubufux/anaconda3/envs/pymc_env/include/python3.11 -I/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/c_code -L/home/ubufux/anaconda3/envs/pymc_env/lib -fvisibility=hidden -o /home/ubufux/.pytensor/compiledir_Linux-4.19-microsoft-standard-x86_64-with-glibc2.31-x86_64-3.11.4-64/tmpfiz1ydb9/m9ba99dab46aeb8bbe4b53a1ff4656102dd2449158b4a891b44897c8c3c627398.so /home/ubufux/.pytensor/compiledir_Linux-4.19-microsoft-standard-x86_64-with-glibc2.31-x86_64-3.11.4-64/tmpfiz1ydb9/mod.cpp -lpython3.11

cc1plus: out of memory allocating 65536 bytes after a total of 7837577216 bytes


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ubufux/./skripte/MCMC_CD.py", line 118, in <module>
    trace = pm.sample(2500, tune=100,nuts_sampler= "numpyro") #nuts_sampler= "numpyro",    nuts_sampler= "blackjax"
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 653, in sample
    step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 233, in assign_step_methods
    return instantiate_steppers(model, steps, selected_steps, step_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 134, in instantiate_steppers
    step = step_class(vars=vars, model=model, **args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/step_methods/hmc/nuts.py", line 180, in __init__
    super().__init__(vars, **kwargs)
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/step_methods/hmc/base_hmc.py", line 109, in __init__
    super().__init__(vars, blocked=blocked, model=self._model, dtype=dtype, **pytensor_kwargs)
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/step_methods/arraystep.py", line 164, in __init__
    func = model.logp_dlogp_function(vars, dtype=dtype, **pytensor_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/model.py", line 606, in logp_dlogp_function
    return ValueGradFunction(costs, grad_vars, extra_vars_and_values, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/model.py", line 345, in __init__
    self._pytensor_function = compile_pymc(inputs, outputs, givens=givens, **kwargs)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/pytensorf.py", line 1196, in compile_pymc
    pytensor_function = pytensor.function(
                        ^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/compile/function/__init__.py", line 315, in function
    fn = pfunc(
         ^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py", line 367, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1756, in orig_function
    fn = m.create(defaults)
         ^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1649, in create
    _fn, _i, _o = self.linker.make_thunk(
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/basic.py", line 254, in make_thunk
    return self.make_all(
           ^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/vm.py", line 1252, in make_all
    raise_with_op(fgraph, node)
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/utils.py", line 535, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/vm.py", line 1243, in make_all
    node.op.make_thunk(node, storage_map, compute_map, [], impl=impl)
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/op.py", line 131, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/op.py", line 96, in make_c_thunk
    outputs = cl.make_thunk(
              ^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1200, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
                                                             ^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1120, in __compile__
    thunk, module = self.cthunk_factory(
                    ^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1644, in cthunk_factory
    module = cache.module_from_key(key=key, lnk=self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 1240, in module_from_key
    module = lnk.compile_cmodule(location)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/basic.py", line 1543, in compile_cmodule
    module = c_compiler.compile_str(
             ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/cmodule.py", line 2649, in compile_str
    raise CompileError(
pytensor.link.c.exceptions.CompileError: Compilation failed (return status=1):
/home/ubufux/anaconda3/envs/pymc_env/bin/g++ -shared -g -O3 -fno-math-errno -Wno-unused-label -Wno-unused-variable -Wno-write-strings -Wno-c++11-narrowing -fno-exceptions -fno-unwind-tables -fno-asynchronous-unwind-tables -march=skylake -mmmx -mpopcnt -msse -msse2 -msse3 -mssse3 -msse4.1 -msse4.2 -mavx -mavx2 -mno-sse4a -mno-fma4 -mno-xop -mfma -mno-avx512f -mbmi -mbmi2 -maes -mpclmul -mno-avx512vl -mno-avx512bw -mno-avx512dq -mno-avx512cd -mno-avx512er -mno-avx512pf -mno-avx512vbmi -mno-avx512ifma -mno-avx5124vnniw -mno-avx5124fmaps -mno-avx512vpopcntdq -mno-avx512vbmi2 -mno-gfni -mno-vpclmulqdq -mno-avx512vnni -mno-avx512bitalg -mno-avx512bf16 -mno-avx512vp2intersect -mno-3dnow -madx -mabm -mno-cldemote -mclflushopt -mno-clwb -mno-clzero -mcx16 -mno-enqcmd -mf16c -mfsgsbase -mfxsr -mno-hle -msahf -mno-lwp -mlzcnt -mmovbe -mno-movdir64b -mno-movdiri -mno-mwaitx -mno-pconfig -mno-pku -mno-prefetchwt1 -mprfchw -mno-ptwrite -mno-rdpid -mrdrnd -mrdseed -mno-rtm -mno-serialize -mno-sgx -mno-sha -mno-shstk -mno-tbm -mno-tsxldtrk -mno-vaes -mno-waitpkg -mno-wbnoinvd -mxsave -mxsavec -mxsaveopt -mxsaves -mno-amx-tile -mno-amx-int8 -mno-amx-bf16 -mno-uintr -mno-hreset -mno-kl -mno-widekl -mno-avxvnni -mno-avx512fp16 --param l1-cache-size=32 --param l1-cache-line-size=64 --param l2-cache-size=8192 -mtune=skylake -DNPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION -m64 -fPIC -I/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/numpy/core/include -I/home/ubufux/anaconda3/envs/pymc_env/include/python3.11 -I/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/link/c/c_code -L/home/ubufux/anaconda3/envs/pymc_env/lib -fvisibility=hidden -o /home/ubufux/.pytensor/compiledir_Linux-4.19-microsoft-standard-x86_64-with-glibc2.31-x86_64-3.11.4-64/tmpfiz1ydb9/m9ba99dab46aeb8bbe4b53a1ff4656102dd2449158b4a891b44897c8c3c627398.so /home/ubufux/.pytensor/compiledir_Linux-4.19-microsoft-standard-x86_64-with-glibc2.31-x86_64-3.11.4-64/tmpfiz1ydb9/mod.cpp -lpython3.11

cc1plus: out of memory allocating 65536 bytes after a total of 7837577216 bytes

Apply node that caused the error: Composite{...}(ExpandDims{axis=0}.0, [ 0.000000 ... 06443e-02], Mul.0, ExpandDims{axis=0}.0, [ 0.000000 ... 04684e-04], [-503.27126321], [20.], Composite{...}.0, [342.], Composite{...}.0, [2754.], Composite{...}.0, [4080.], Composite{(sqr(sqr(i0)) * i0)}.0, [21840.], Composite{(sqr(sqr(i0)) * sqr(i0))}.0, [4.], Composite{(i2 - (i0 * i1))}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, [  0. ... .        ], [19.], ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, [-0.82608696], [95.], [108.], [306.], [17.], [102.], [272.], [32.], [96.], [240.], [45.], [90.], [210.], [56.], [84.], [182.], [65.], [78.], [156.], [72.], [132.], [77.], [66.], [110.], [80.], [60.], [81.], [54.], [48.], [42.], [36.], [30.], [24.], [12.], [18.], [6.], [2.], [1530.], [5508.], [1632.], [4624.], [1440.], [3840.], [420.], [1260.], [3150.], [546.], [1092.], [2548.], [624.], [936.], [2028.], [660.], [792.], [1584.], [1210.], [630.], [540.], [900.], [576.], [432.], [648.], [504.], [336.], [448.], [252.], [294.], [330.], [180.], [120.], [100.], [11560.], [41616.], [1088.], [2040.], [11520.], [32640.], [2160.], [3360.], [1575.], [9450.], [25200.], [4095.], [7644.], [19110.], [3822.], [4368.], [3042.], [6084.], [14196.], [390.], [4290.], [3168.], [4752.], [10296.], [4356.], [3960.], [3025.], [3630.], [7260.], [1058.75], [4331.25], [3465.], [2700.], [4950.], [1350.], [4050.], [2880.], [2268.], [1944.], [3240.], [1490.4], [3693.6], [1792.], [1344.], [2016.], [1568.], [3192.], [1680.], [1323.], [882.], [1176.], [1501.5], [2656.5], [1155.], [2070.], [720.], [1095.71428571], [1504.28571429], [550.], [300.], [972.], [168.], [288.], [144.], [466.875], [523.125], [117.], [28.], [186.], [16800.], [60480.], [13440.], [16380.], [46410.], [24570.], [38220.], [2184.], [13104.], [34944.], [32760.], [42588.], [1716.], [36036.], [41184.], [3432.], [25740.], [3300.], [36300.], [7920.], [18480.], [5940.], [32670.], [29700.], [12870.], [6930.], [28350.], [22680.], [3600.], [4320.], [8640.], [7560.], [16128.], [3024.], [5544.], [6955.2], [17236.8], [10584.], [2352.], [5880.], [11970.], [6300.], [1890.], [7590.], [1080.], [960.], [4140.], [600.], [360.], [1314.85714286], [1805.14285714], [468.], [264.], [396.], [486.], [81900.], [294840.], [65520.], [122850.], [72072.], [204204.], [108108.], [168168.], [8580.], [51480.], [137280.], [128700.], [167310.], [124740.], [142560.], [11880.], [35640.], [89100.], [9900.], [108900.], [23760.], [55440.], [15120.], [83160.], [75600.], [10080.], [14553.], [59535.], [47628.], [9072.], [18144.], [12600.], [37800.], [26880.], [5040.], [9240.], [8694.], [21546.], [13230.], [2940.], [2520.], [4200.], [10260.], [5400.], [1620.], [2145.], [3795.], [1650.], [480.], [828.], [23.], [506.], ExpandDims{axis=0}.0, [5313.], [8740.], [66861.], ExpandDims{axis=0}.0, obs{[-24454.60 ... .38516757]}, [1.], Composite{...}.1, Composite{...}.1, Composite{...}.1, Mul.0, [16.], [15.], [14.], [13.], [11.], [10.], [9.], [8.], [7.], [5.], [3.], [2448.], [1638.], [1404.], [1188.], [990.], [810.], [378.], [270.], [5460.], [5148.], [4620.], [1848.], [780.], [24024.], [21450.], [17820.], [13860.], [6804.], [2310.], [0.04347826], [-13.30434783], [-106.43478261], [-219.13043478], [559.56521739], [3286.95652174], [3.39130435], [34.43478261], [131.52173913], [157.82608696], [46.0326087], [188.31521739], [301.30434783], [1232.60869565], [156.52173913], [187.82608696], [632.73913043], [2588.47826087], [394.43478261], [6.7826087], [68.86956522], [315.65217391], [150.65217391], [986.08695652], [375.65217391], [2070.7826087], [788.86956522], [28.69565217], [117.39130435], [176.08695652], [131.47826087], [547.82608696], [1643.47826087], [219.13043478], [0.04347826], [5.73913043], [52.60869565], [215.2173913], [125.2173913], [701.2173913], [241.04347826], [1168.69565217], [401.73913043], [2.86956522], [27.39130435], [98.60869565], [84.52173913], [64.8], [160.59130435], [302.4], [749.42608696], [102.26086957], [87.65217391], [936.7826087], [127.82608696], [109.56521739], [0.04347826], [4.7826087], [39.13043478], [140.86956522], [460.17391304], [575.2173913], [182.60869565], [2.60869565], [25.04347826], [18.7826087], [77.91304348], [58.43478261], [68.17391304], [138.7826087], [255.65217391], [520.43478261], [73.04347826], [446.08695652], [46.95652174], [0.04347826], [28.17391304], [273.91304348], [82.17391304], [234.7826087], [70.43478261], [3.52173913], [21.91304348], [14.60869565], [57.52173913], [38.34782609], [65.2826087], [115.5], [31.30434783], [93.26086957], [165.], [0.04347826], [0.04347826], [3.13043478], [19.47826087], [51.13043478], [50.2173913], [41.73913043], [71.73913043], [20.86956522], [2.08695652], [10.95652174], [26.08695652], [5.2173913], [0.04347826], [0.04347826], [0.04347826], [2.43478261], [12.7826087], [18.26086957], [12.52173913], [1.82608696], [14.34782609], [7.82608696], [47.63975155], [65.40372671], [23.91304348], [13.04347826], [57.16770186], [78.48447205], [11.47826087], [0.04347826], [20.34782609], [6.26086957], [42.26086957], [17.2173913], [21.13043478], [0.04347826], [0.04347826], [0.04347826], [4.34782609], [7.30434783], [1.30434783], [20.29891304], [22.74456522], [5.08695652], [0.04347826], [0.04347826], [1.95652174], [1.56521739], [1.04347826], [1.2173913], [8.08695652], [0.04347826], [0.04347826], [0.52173913], [0.7826087], [0.04347826], [0.17391304], [0.04347826], [0.08695652], [0.26086957], [22.], [231.], [380.], [2907.], ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, ExpandDims{axis=0}.0, [-0.5], [0.91893853], Log.0, Gt.0, [-inf])
Toposort index: 21
Inputs types: [TensorType(float64, shape=(1,)), TensorType(float64, shape=(100,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(100,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(100,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(float64, shape=(1,)), TensorType(bool, shape=(1,)), TensorType(float32, shape=(1,))]

HINT: Use a linker other than the C linker to print the inputs' shapes and strides.
HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

How large is your dataset?

There are only 100 points for now (picture).

A little context… X-axis is temperature. Parameter β€œw” in the model is a vector of ws for each temperature, so β€œcomplicated” polynomial P(v,w) has dimension of X (len=100) at end.

image

Any chance you can provide a reproducible example and maybe figure out at how many terms in the polynomial do things break down?

Sure here it is. For N_sim = 7 it works well and fast enough. For N_sim=17 it is already slow. And for N=22 I encounter compilation error posted in previous post.

# -*- coding: utf-8 -*-
"""
Created on Wed Jul 19 10:42:25 2023

@author: urosz
"""
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 27 13:05:53 2023

@author: urosz
"""
import pymc as pm
import os
import pytensor as pt
import numpy as np
import matplotlib.pyplot as plt
import arviz as az
import pymc.sampling_jax
import timeit

T0=273.15; R=1.987*10**(-3)

def LR_functions(w,v,H1,H2,C,N):
    if N==7:
        Q=3*v**4*w**2 + 9*v**4*w + 12*v**4 + 2*v**3*w**3 + 6*v**3*w**2 + 12*v**3*w + 20*v**3 + v**2*w**5 + 2*v**2*w**4 + 3*v**2*w**3 + 4*v**2*w**2 + 5*v**2*w + 21*v**2 + 7*v + 1
        cd=96*C*v**4 + 160*C*v**3 + 168*C*v**2 + 56*C*v + 8*C + 8*H1*v**4*w**2 + 2*v**4*w**2*(3*C + 5*H1) + 36*v**4*w*(C + H1) + 4*v**3*w**3*(C + 3*H1) + 6*v**3*w**2*(3*C + 5*H1) + 48*v**3*w*(C + H1) + 2*v**2*w**5*(3*H1 + H2) + 2*v**2*w**4*(C + 6*H1 + H2) + 6*v**2*w**3*(C + 3*H1) + 4*v**2*w**2*(3*C + 5*H1) + 20*v**2*w*(C + H1)
                 
    elif N== 17:
        Q=54*v**6*w**8 + 240*v**6*w**7 + 630*v**6*w**6 + 1260*v**6*w**5 + 2100*v**6*w**4 + 3024*v**6*w**3 + 3780*v**6*w**2 + 3960*v**6*w + 33*v**5*w**10 + 120*v**5*w**9 + 270*v**5*w**8 + 480*v**5*w**7 + 735*v**5*w**6 + 1008*v**5*w**5 + 1260*v**5*w**4 + 1440*v**5*w**3 + 1485*v**5*w**2 + 1320*v**5*w + 13*v**4*w**12 + 39*v**4*w**11 + 78*v**4*w**10 + 130*v**4*w**9 + 195*v**4*w**8 + 273*v**4*w**7 + 364*v**4*w**6 + 468*v**4*w**5 + 585*v**4*w**4 + 715*v**4*w**3 + 858*v**4*w**2 + 1014*v**4*w + 1092*v**4 + 2*v**3*w**13 + 6*v**3*w**12 + 12*v**3*w**11 + 20*v**3*w**10 + 30*v**3*w**9 + 42*v**3*w**8 + 56*v**3*w**7 + 72*v**3*w**6 + 90*v**3*w**5 + 110*v**3*w**4 + 132*v**3*w**3 + 156*v**3*w**2 + 182*v**3*w + 210*v**3 + v**2*w**15 + 2*v**2*w**14 + 3*v**2*w**13 + 4*v**2*w**12 + 5*v**2*w**11 + 6*v**2*w**10 + 7*v**2*w**9 + 8*v**2*w**8 + 9*v**2*w**7 + 10*v**2*w**6 + 11*v**2*w**5 + 12*v**2*w**4 + 13*v**2*w**3 + 14*v**2*w**2 + 15*v**2*w + 136*v**2 + 17*v + 1
        cd=19656*C*v**4 + 3780*C*v**3 + 2448*C*v**2 + 306*C*v + 18*C + v**6*w**8*(84*C + 72*H1 + 60*H2) + v**6*w**8*(168.0*C + 472.5*H1 + 115.5*H2) + v**6*w**7*(480*C + 360*H1 + 240*H2) + v**6*w**7*(900.0*C + 1980.0*H1 + 360.0*H2) + v**6*w**6*(1620*C + 1080*H1 + 540*H2) + v**6*w**6*(2700.0*C + 4950.0*H1 + 450.0*H2) + v**6*w**5*(4200*C + 2520*H1 + 840*H2) + v**6*w**5*(5880.0*C + 8820.0*H1 + 420.0*H2) + v**6*w**4*(10080*C + 12600*H1) + v**6*w**4*(9240*C + 5040*H1 + 840*H2) + v**6*w**3*(13608*C + 13608*H1) + v**6*w**3*(18144*C + 9072*H1) + v**6*w**2*(12600*C + 10080*H1) + v**6*w**2*(32760*C + 12600*H1) + 7920*v**6*w*(7*C + 2*H1) + v**5*w**10*(30*C + 36*H1 + 42*H2) + v**5*w**10*(54.0*C + 307.8*H1 + 124.2*H2) + v**5*w**9*(144*C + 144*H1 + 144*H2) + v**5*w**9*(288.0*C + 1080.0*H1 + 360.0*H2) + v**5*w**8*(420*C + 360*H1 + 300*H2) + v**5*w**8*(840.0*C + 2362.5*H1 + 577.5*H2) + v**5*w**7*(960*C + 720*H1 + 480*H2) + v**5*w**7*(1800.0*C + 3960.0*H1 + 720.0*H2) + v**5*w**6*(1890*C + 1260*H1 + 630*H2) + v**5*w**6*(3150.0*C + 5775.0*H1 + 525.0*H2) + v**5*w**5*(3360*C + 2016*H1 + 672*H2) + v**5*w**5*(4704.0*C + 7056.0*H1 + 336.0*H2) + v**5*w**4*(6048*C + 7560*H1) + v**5*w**4*(5544*C + 3024*H1 + 504*H2) + v**5*w**3*(6480*C + 6480*H1) + v**5*w**3*(8640*C + 4320*H1) + v**5*w**2*(4950*C + 3960*H1) + v**5*w**2*(12870*C + 4950*H1) + 2640*v**5*w*(7*C + 2*H1) + v**4*w**12*(126.5*H1 + 71.5*H2) + v**4*w**12*(6*C + 12*H1 + 18*H2) + v**4*w**11*(30.0*C + 342.0*H1 + 168.0*H2) + v**4*w**11*(36*C + 54*H1 + 72*H2) + v**4*w**10*(108.0*C + 615.6*H1 + 248.4*H2) + v**4*w**10*(120*C + 144*H1 + 168*H2) + v**4*w**9*(240.0*C + 900.0*H1 + 300.0*H2) + v**4*w**9*(300*C + 300*H1 + 300*H2) + v**4*w**8*(420.0*C + 1181.25*H1 + 288.75*H2) + v**4*w**8*(630*C + 540*H1 + 450*H2) + v**4*w**7*(630.0*C + 1386.0*H1 + 252.0*H2) + v**4*w**7*(1176*C + 882*H1 + 588*H2) + v**4*w**6*(840.0*C + 1540.0*H1 + 140.0*H2) + v**4*w**6*(2016*C + 1344*H1 + 672*H2) + v**4*w**5*(1008.0*C + 1512.0*H1 + 72.0*H2) + v**4*w**5*(3240*C + 1944*H1 + 648*H2) + v**4*w**4*(1080*C + 1350*H1) + v**4*w**4*(4950*C + 2700*H1 + 450*H2) + v**4*w**3*(990*C + 990*H1) + v**4*w**3*(7260*C + 3630*H1) + v**4*w**2*(660*C + 528*H1) + v**4*w**2*(10296*C + 3960*H1) + 2028*v**4*w*(7*C + 2*H1) + v**3*w**13*(4*C + 12*H1 + 20*H2) + v**3*w**12*(18*C + 36*H1 + 54*H2) + v**3*w**11*(48*C + 72*H1 + 96*H2) + v**3*w**10*(100*C + 120*H1 + 140*H2) + v**3*w**9*(180*C + 180*H1 + 180*H2) + v**3*w**8*(294*C + 252*H1 + 210*H2) + v**3*w**7*(448*C + 336*H1 + 224*H2) + v**3*w**6*(648*C + 432*H1 + 216*H2) + v**3*w**5*(900*C + 540*H1 + 180*H2) + v**3*w**4*(1210*C + 660*H1 + 110*H2) + v**3*w**3*(1584*C + 792*H1) + v**3*w**2*(2028*C + 780*H1) + 364*v**3*w*(7*C + 2*H1) + v**2*w**15*(6*H1 + 12*H2) + v**2*w**14*(2*C + 12*H1 + 22*H2) + v**2*w**13*(6*C + 18*H1 + 30*H2) + v**2*w**12*(12*C + 24*H1 + 36*H2) + v**2*w**11*(20*C + 30*H1 + 40*H2) + v**2*w**10*(30*C + 36*H1 + 42*H2) + v**2*w**9*(42*C + 42*H1 + 42*H2) + v**2*w**8*(56*C + 48*H1 + 40*H2) + v**2*w**7*(72*C + 54*H1 + 36*H2) + v**2*w**6*(90*C + 60*H1 + 30*H2) + v**2*w**5*(110*C + 66*H1 + 22*H2) + v**2*w**4*(132*C + 72*H1 + 12*H2) + v**2*w**3*(156*C + 78*H1) + v**2*w**2*(182*C + 70*H1) + 30*v**2*w*(7*C + 2*H1)
    
    elif N==22:
        Q=84*v**6*w**13 + 390*v**6*w**12 + 1080*v**6*w**11 + 2310*v**6*w**10 + 4200*v**6*w**9 + 6804*v**6*w**8 + 10080*v**6*w**7 + 13860*v**6*w**6 + 17820*v**6*w**5 + 21450*v**6*w**4 + 24024*v**6*w**3 + 24570*v**6*w**2 + 21840*v**6*w + 48*v**5*w**15 + 180*v**5*w**14 + 420*v**5*w**13 + 780*v**5*w**12 + 1260*v**5*w**11 + 1848*v**5*w**10 + 2520*v**5*w**9 + 3240*v**5*w**8 + 3960*v**5*w**7 + 4620*v**5*w**6 + 5148*v**5*w**5 + 5460*v**5*w**4 + 5460*v**5*w**3 + 5040*v**5*w**2 + 4080*v**5*w + 18*v**4*w**17 + 54*v**4*w**16 + 108*v**4*w**15 + 180*v**4*w**14 + 270*v**4*w**13 + 378*v**4*w**12 + 504*v**4*w**11 + 648*v**4*w**10 + 810*v**4*w**9 + 990*v**4*w**8 + 1188*v**4*w**7 + 1404*v**4*w**6 + 1638*v**4*w**5 + 1890*v**4*w**4 + 2160*v**4*w**3 + 2448*v**4*w**2 + 2754*v**4*w + 2907*v**4 + 2*v**3*w**18 + 6*v**3*w**17 + 12*v**3*w**16 + 20*v**3*w**15 + 30*v**3*w**14 + 42*v**3*w**13 + 56*v**3*w**12 + 72*v**3*w**11 + 90*v**3*w**10 + 110*v**3*w**9 + 132*v**3*w**8 + 156*v**3*w**7 + 182*v**3*w**6 + 210*v**3*w**5 + 240*v**3*w**4 + 272*v**3*w**3 + 306*v**3*w**2 + 342*v**3*w + 380*v**3 + v**2*w**20 + 2*v**2*w**19 + 3*v**2*w**18 + 4*v**2*w**17 + 5*v**2*w**16 + 6*v**2*w**15 + 7*v**2*w**14 + 8*v**2*w**13 + 9*v**2*w**12 + 10*v**2*w**11 + 11*v**2*w**10 + 12*v**2*w**9 + 13*v**2*w**8 + 14*v**2*w**7 + 15*v**2*w**6 + 16*v**2*w**5 + 17*v**2*w**4 + 18*v**2*w**3 + 19*v**2*w**2 + 20*v**2*w + 231*v**2 + 22*v + 1
        cd=66861*C*v**4 + 8740*C*v**3 + 5313*C*v**2 + 506*C*v + 23*C + v**6*w**13*(84*C + 72*H1 + 120*H2) + v**6*w**13*(288.0*C + 828.0*H1 + 540.0*H2) + v**6*w**12*(480*C + 360*H1 + 540*H2) + v**6*w**12*(1650.0*C + 3795.0*H1 + 2145.0*H2) + v**6*w**11*(1620*C + 1080*H1 + 1440*H2) + v**6*w**11*(5400.0*C + 10260.0*H1 + 5040.0*H2) + v**6*w**10*(4200*C + 2520*H1 + 2940*H2) + v**6*w**10*(13230.0*C + 21546.0*H1 + 8694.0*H2) + v**6*w**9*(9240*C + 5040*H1 + 5040*H2) + v**6*w**9*(26880.0*C + 37800.0*H1 + 12600.0*H2) + v**6*w**8*(18144*C + 9072*H1 + 7560*H2) + v**6*w**8*(47628.0*C + 59535.0*H1 + 14553.0*H2) + v**6*w**7*(32760*C + 15120*H1 + 10080*H2) + v**6*w**7*(75600.0*C + 83160.0*H1 + 15120.0*H2) + v**6*w**6*(55440*C + 23760*H1 + 11880*H2) + v**6*w**6*(108900.0*C + 108900.0*H1 + 9900.0*H2) + v**6*w**5*(89100*C + 35640*H1 + 11880*H2) + v**6*w**5*(142560.0*C + 124740.0*H1 + 5940.0*H2) + v**6*w**4*(167310*C + 128700*H1) + v**6*w**4*(137280*C + 51480*H1 + 8580*H2) + v**6*w**3*(168168*C + 108108*H1) + v**6*w**3*(204204*C + 72072*H1) + v**6*w**2*(122850*C + 65520*H1) + v**6*w**2*(294840*C + 81900*H1) + 21840*v**6*w*(19*C + 4*H1) + v**5*w**15*(30*C + 36*H1 + 72*H2) + v**5*w**15*(84.0*C + 486.0*H1 + 396.0*H2) + v**5*w**14*(144*C + 144*H1 + 264*H2) + v**5*w**14*(468.0*C + 1805.14285714286*H1 + 1314.85714285714*H2) + v**5*w**13*(420*C + 360*H1 + 600*H2) + v**5*w**13*(1440.0*C + 4140.0*H1 + 2700.0*H2) + v**5*w**12*(960*C + 720*H1 + 1080*H2) + v**5*w**12*(3300.0*C + 7590.0*H1 + 4290.0*H2) + v**5*w**11*(1890*C + 1260*H1 + 1680*H2) + v**5*w**11*(6300.0*C + 11970.0*H1 + 5880.0*H2) + v**5*w**10*(3360*C + 2016*H1 + 2352*H2) + v**5*w**10*(10584.0*C + 17236.8*H1 + 6955.2*H2) + v**5*w**9*(5544*C + 3024*H1 + 3024*H2) + v**5*w**9*(16128.0*C + 22680.0*H1 + 7560.0*H2) + v**5*w**8*(8640*C + 4320*H1 + 3600*H2) + v**5*w**8*(22680.0*C + 28350.0*H1 + 6930.0*H2) + v**5*w**7*(12870*C + 5940*H1 + 3960*H2) + v**5*w**7*(29700.0*C + 32670.0*H1 + 5940.0*H2) + v**5*w**6*(18480*C + 7920*H1 + 3960*H2) + v**5*w**6*(36300.0*C + 36300.0*H1 + 3300.0*H2) + v**5*w**5*(25740*C + 10296*H1 + 3432*H2) + v**5*w**5*(41184.0*C + 36036.0*H1 + 1716.0*H2) + v**5*w**4*(42588*C + 32760*H1) + v**5*w**4*(34944*C + 13104*H1 + 2184*H2) + v**5*w**3*(38220*C + 24570*H1) + v**5*w**3*(46410*C + 16380*H1) + v**5*w**2*(25200*C + 13440*H1) + v**5*w**2*(60480*C + 16800*H1) + 4080*v**5*w*(19*C + 4*H1) + v**4*w**17*(186.0*H1 + 182.0*H2) + v**4*w**17*(6*C + 12*H1 + 28*H2) + v**4*w**16*(36*C + 54*H1 + 117*H2) + v**4*w**16*(45.0*C + 523.125*H1 + 466.875*H2) + v**4*w**15*(120*C + 144*H1 + 288*H2) + v**4*w**15*(168.0*C + 972.0*H1 + 792.0*H2) + v**4*w**14*(300*C + 300*H1 + 550*H2) + v**4*w**14*(390.0*C + 1504.28571428571*H1 + 1095.71428571429*H2) + v**4*w**13*(630*C + 540*H1 + 900*H2) + v**4*w**13*(720.0*C + 2070.0*H1 + 1350.0*H2) + v**4*w**12*(1155.0*C + 2656.5*H1 + 1501.5*H2) + v**4*w**12*(1176*C + 882*H1 + 1323*H2) + v**4*w**11*(1680.0*C + 3192.0*H1 + 1568.0*H2) + v**4*w**11*(2016*C + 1344*H1 + 1792*H2) + v**4*w**10*(2268.0*C + 3693.6*H1 + 1490.4*H2) + v**4*w**10*(3240*C + 1944*H1 + 2268*H2) + v**4*w**9*(2880.0*C + 4050.0*H1 + 1350.0*H2) + v**4*w**9*(4950*C + 2700*H1 + 2700*H2) + v**4*w**8*(3465.0*C + 4331.25*H1 + 1058.75*H2) + v**4*w**8*(7260*C + 3630*H1 + 3025*H2) + v**4*w**7*(3960.0*C + 4356.0*H1 + 792.0*H2) + v**4*w**7*(10296*C + 4752*H1 + 3168*H2) + v**4*w**6*(4290.0*C + 4290.0*H1 + 390.0*H2) + v**4*w**6*(14196*C + 6084*H1 + 3042*H2) + v**4*w**5*(4368.0*C + 3822.0*H1 + 182.0*H2) + v**4*w**5*(19110*C + 7644*H1 + 2548*H2) + v**4*w**4*(4095*C + 3150*H1) + v**4*w**4*(25200*C + 9450*H1 + 1575*H2) + v**4*w**3*(3360*C + 2160*H1) + v**4*w**3*(32640*C + 11520*H1) + v**4*w**2*(2040*C + 1088*H1) + v**4*w**2*(41616*C + 11560*H1) + 2754*v**4*w*(19*C + 4*H1) + v**3*w**18*(4*C + 12*H1 + 30*H2) + v**3*w**17*(18*C + 36*H1 + 84*H2) + v**3*w**16*(48*C + 72*H1 + 156*H2) + v**3*w**15*(100*C + 120*H1 + 240*H2) + v**3*w**14*(180*C + 180*H1 + 330*H2) + v**3*w**13*(294*C + 252*H1 + 420*H2) + v**3*w**12*(448*C + 336*H1 + 504*H2) + v**3*w**11*(648*C + 432*H1 + 576*H2) + v**3*w**10*(900*C + 540*H1 + 630*H2) + v**3*w**9*(1210*C + 660*H1 + 660*H2) + v**3*w**8*(1584*C + 792*H1 + 660*H2) + v**3*w**7*(2028*C + 936*H1 + 624*H2) + v**3*w**6*(2548*C + 1092*H1 + 546*H2) + v**3*w**5*(3150*C + 1260*H1 + 420*H2) + v**3*w**4*(3840*C + 1440*H1 + 240*H2) + v**3*w**3*(4624*C + 1632*H1) + v**3*w**2*(5508*C + 1530*H1) + 342*v**3*w*(19*C + 4*H1) + v**2*w**20*(6*H1 + 17*H2) + v**2*w**19*(2*C + 12*H1 + 32*H2) + v**2*w**18*(6*C + 18*H1 + 45*H2) + v**2*w**17*(12*C + 24*H1 + 56*H2) + v**2*w**16*(20*C + 30*H1 + 65*H2) + v**2*w**15*(30*C + 36*H1 + 72*H2) + v**2*w**14*(42*C + 42*H1 + 77*H2) + v**2*w**13*(56*C + 48*H1 + 80*H2) + v**2*w**12*(72*C + 54*H1 + 81*H2) + v**2*w**11*(90*C + 60*H1 + 80*H2) + v**2*w**10*(110*C + 66*H1 + 77*H2) + v**2*w**9*(132*C + 72*H1 + 72*H2) + v**2*w**8*(156*C + 78*H1 + 65*H2) + v**2*w**7*(182*C + 84*H1 + 56*H2) + v**2*w**6*(210*C + 90*H1 + 45*H2) + v**2*w**5*(240*C + 96*H1 + 32*H2) + v**2*w**4*(272*C + 102*H1 + 17*H2) + v**2*w**3*(306*C + 108*H1) + v**2*w**2*(342*C + 95*H1) + 20*v**2*w*(19*C + 4*H1)

    return (Q,cd)

def w_param(dG,dH,dcp,tcd):
    tCD=tcd+273.15
    return np.exp((-1 / R) * (dG/ T0 + dH * ((1 / tCD) - (1 / T0)) + dcp * (1 - (T0 / tCD) - np.log(tCD/ T0))))

    
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

N_sim=7 #dataset

#Parameters for simulating the data "true_params"
dcp_true=-0.008;dH_true= -1.0;dG_true = -0.23
v_true=0.048

dh2_true=100;h2_true=-40000
c_true=2200;ct_true=-30
k_true=4

# # Generate some synthetic data from the model
N = 100
T_CD = np.linspace(0, 100, N)
w_true = w_param(dG_true, dH_true, dcp_true, T_CD)
h2_t = h2_true + dh2_true*T_CD
h1_t = h2_t*(1-k_true/6)
c_t = c_true + ct_true*T_CD

#### RETRIVE FUNCTIONS FROM LR_functions
Q_sim,cd_sim=LR_functions(w_true, v_true, h1_t, h2_t, c_t, N_sim)

# simulated model
Y = (1/(N_sim+1))*(cd_sim/Q_sim) #model function
#add some noise
Y = np.random.normal(Y, 100)

#PLOT simulated data
plt.figure(figsize=(7, 5))
plt.scatter(T_CD, Y, label='Simulated data')
plt.show()

T_CD_K=T_CD + 273.15


with pm.Model() as model:
    #VARIABLES
    dg = pm.Normal('dg', mu=-0.20, sigma=0.2)
    dh = pm.Normal('dh', mu=-1.0, sigma=0.2)
    dcp = pm.Normal('dcp', mu=0, sigma=0.01)
    v = pm.Normal('v', mu=0.06, sigma=0.02)
    H2=pm.Normal('H2', mu=-40000, sigma=5000)
    dH2t=pm.Normal('dH2', mu=1000, sigma=300)
    C=pm.Normal('C', mu=2200, sigma=500)
    C_temp=pm.Normal('C_temp', mu=-50, sigma=50)
    k=pm.Normal('k', mu=3.8, sigma=1)
    
    #Determenistic variables
    w = np.exp((-1 / R) * (dg / T0 + dh * ((1 / T_CD_K) - (1 / T0)) + dcp * (1 - (T0 / T_CD_K) - np.log(T_CD_K/ T0))))
    H2_t=H2 + dH2t*T_CD
    H1_t=H2_t*(1-k/6)
    C_t=C+C_temp*T_CD
    
    
    Q,cd=LR_functions(w,v,H1_t,H2_t,C_t,N_sim) #call functions to calculate Q,cd (polynomials)

    mu = cd/(Q*(N_sim+1)) #final model
    sd = pm.HalfNormal('sd', sigma=300)
    obs = pm.Normal('obs', mu=mu, sigma=sd, observed=Y)

    # Perform MCMC
    trace = pm.sample(2500, tune=100,nuts_sampler= "numpyro") #nuts_sampler= "numpyro",    nuts_sampler= "blackjax"


print(az.summary(trace,var_names=['dg', 'dh','dcp','k', 'H2','sd'], round_to=2))

az.plot_trace(trace, var_names=['dg', 'dh','dcp','k', 'H2','sd'])
plt.show()

#energy plots
az.plot_energy(trace)
plt.show()

with model:
    ppc=pm.sample_posterior_predictive(trace, extend_inferencedata=True, random_seed=42)

trace.posterior_predictive

az.plot_posterior(ppc, var_names=['dg', 'dh','dcp','k', 'H2','sd'])
plt.show()

az.plot_ppc(trace, num_pp_samples=500)
plt.show()

All of those long complex expressions can be gathered into shorter vectorized operations. The Q term can be written:

v_powers = np.array([...])
w_powers = np.array([...])
Q_coefs = np.array([...])
vec_powers = pt.outer(pt.power(v, v_powers),
                      pt.power(w, w_powers)).ravel()
Q = (vec_powers * Q_coefs).sum()

The cd term can be written similary, but you have cd_coefs = a + X @ CH_vec, where X is a (n, 3) matrix of coefficients, a is mostly zeros, and CH_vec is just np.array([C, H1, H2]).

I’m not sure it will help with speed, but at bare minimum it’s a bit easier to read/handle. I have a suspicion it will help though, because polynomials are hard to break apart and simplify. This formulation makes the β€œsteps” that go into their construction really obvious to the compiler.

2 Likes

Sorry for silly question, I am quite in PyMC/PyTensor and I am struggling with this vectorization step. Problem is that β€œw” is actually vector/array ( lenght of N, temperatures). How should I make it work for both β€œv” which is scalar and also β€œw” (vector)?

Traceback (most recent call last):
  File "/home/ubufux/./skripte/MCMC_CD.py", line 73, in <module>
    Q_sim,cd_sim=LR_functions(w_true, v_true, h1_t, h2_t, c_t, N_sim)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/./skripte/MCMC_CD.py", line 31, in LR_functions
    vec_powers = pt.tensor.outer(pt.tensor.power(v, v_powers),pt.tensor.power(w, w_powers)).ravel()
                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubufux/anaconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/tensor/math.py", line 2821, in power
    return x**y
           ~^^~
ValueError: operands could not be broadcast together with shapes (100,) (15,)

You’re just doing a lot of broadcasting, so you need to grok how the dimensions of all your objects work. It’s instructive to work through it in numpy, because pretty much the whole python scientific stack follows the numpy broadcasting framework. You should read about that here.

Here’s some dummy data to work through your specific problem:

w = np.linspace(0, 1, 5)
v = 3
w_powers = np.arange(10)
v_powers = np.arange(4)

so w has shape (5,) and v is a scalar, w_powers and v_powers have shapes (10,) and (4,) respectively.

First, np.power will broadcast exponentiation the same as addition or multiplication. v is trivial, and results in a (4,) array:

v_pow = np.power(v, v_powers)
v_pow
>>> array([ 1,  3,  9, 27], dtype=int32)

Now if you try np.power(w, w_powers) you will get an error, because numpy doesn’t know how to put two objects with shapes (5,) and (10,) together. As an aside, if w_powers were np.arange(5) you would get back the elementwise powers with shape (5,), but that’s not what you want. You want the same set of powers applied to each different w, so the result is a matrix (w_size, w_power_size). So you need to introduce a batch dimension, which can be done by indexing w with None:

w_pow = np.power(w[:, None], w_powers)
w_pow
>>> array([[1.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
>>>       [1.  , 0.25, 0.06, 0.02, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
>>>       [1.  , 0.5 , 0.25, 0.12, 0.06, 0.03, 0.02, 0.01, 0.  , 0.  ],
>>>       [1.  , 0.75, 0.56, 0.42, 0.32, 0.24, 0.18, 0.13, 0.1 , 0.08],
>>>       [1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  ]])

As you can see, the rows are the w values and the columns are the powers, so the (i, j)th element is w[i] ** w_power[j]. The output shape is (5,10), because we started with a (5, 1) array w[:, None], and applied w_power[i] to it for each element in w_power.

So moving on, we now have a (4,) array and a (5, 10) and of course theses don’t conform, so we have to add more batch dimensions using None. The desired output shape now is (4, 5 ,10), because that will have, at position (i, j, k) the product v ** v_power[i] * w[j] ** w_power[k] So insert batch dimensions into v_pow and multiply. Note that can’t use the vector-vector outer product anymore, because now you have a matrix and a vector. If you try, numpy will quietly ravel the matrix and you’ll get the wrong answer.

vw = v_pow[:, None, None] * w_pow
vw 
>>>array([[[ 1.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [ 1.  ,  0.25,  0.06,  0.02,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [ 1.  ,  0.5 ,  0.25,  0.12,  0.06,  0.03,  0.02,  0.01,  0.  ,  0.  ],
>>>        [ 1.  ,  0.75,  0.56,  0.42,  0.32,  0.24,  0.18,  0.13,  0.1 ,  0.08],
>>>        [ 1.  ,  1.  ,  1.  ,  1.  ,  1.  ,  1.  ,  1.  ,  1.  ,  1.  ,  1.  ]],
>>>
>>>       [[ 3.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [ 3.  ,  0.75,  0.19,  0.05,  0.01,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [ 3.  ,  1.5 ,  0.75,  0.38,  0.19,  0.09,  0.05,  0.02,  0.01,  0.01],
>>>        [ 3.  ,  2.25,  1.69,  1.27,  0.95,  0.71,  0.53,  0.4 ,  0.3 ,  0.23],
>>>        [ 3.  ,  3.  ,  3.  ,  3.  ,  3.  ,  3.  ,  3.  ,  3.  ,  3.  ,  3.  ]],
>>>
>>>       [[ 9.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [ 9.  ,  2.25,  0.56,  0.14,  0.04,  0.01,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [ 9.  ,  4.5 ,  2.25,  1.12,  0.56,  0.28,  0.14,  0.07,  0.04,  0.02],
>>>        [ 9.  ,  6.75,  5.06,  3.8 ,  2.85,  2.14,  1.6 ,  1.2 ,  0.9 ,  0.68],
>>>        [ 9.  ,  9.  ,  9.  ,  9.  ,  9.  ,  9.  ,  9.  ,  9.  ,  9.  ,  9.  ]],
>>>
>>>       [[27.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
>>>        [27.  ,  6.75,  1.69,  0.42,  0.11,  0.03,  0.01,  0.  ,  0.  ,  0.  ],
>>>        [27.  , 13.5 ,  6.75,  3.38,  1.69,  0.84,  0.42,  0.21,  0.11,  0.05],
>>>        [27.  , 20.25, 15.19, 11.39,  8.54,  6.41,  4.81,  3.6 ,  2.7 ,  2.03],
>>>        [27.  , 27.  , 27.  , 27.  , 27.  , 27.  , 27.  , 27.  , 27.  , 27.  ]]])

You should verify that the value in a randomly chosen cell is the product I claimed it is.

So then you’ll have to take this object, potentially reshape it, then multiply by the Q coefficients. If Q_coefs has shape (40,) (that is, w_powers.shape[0] * v_powers.shape[0] – one coefficient for each polynomial term) you can reshape the output vw so that each row is np.outer(w[i] ** w_powers, v ** v_powers).ravel(), which would be vw.transpose(1, 2, 0).reshape(5, -1). How did I arrive at the order (1, 2, 0)? All numpy operations work on dimensions left-to-right, so I want 5 rows, I put the size 5 dimension first (the w dimension). Then I want things to come out the same way as that outer product, which will be (w_powers, v_powers), so that’s dim 2 then dim 0 in vw. Then you could broadcast Q_coefs across all 5 w values, and take the sum, which would give you the polynomials evaluated at each of the values of w.

Of course, you have to pay attention to how Q_coefs is arranged for how to reshape vw in the end.

2 Likes

I implemented my model function via numpy array manipultaion as suggested. It works fine, but how do I now translate this procedure into pytensor function so I can use it in pymc model?

import numpy as np
import re
import matplotlib.pyplot as plt

def extract_powers(term):
    """Extract powers of v and w from a term."""
    v_power, w_power = 0, 0
    v_match = re.search(r'v\*\*(\d+)', term)
    w_match = re.search(r'w\*\*(\d+)', term)
    if v_match:
        v_power = int(v_match.group(1))
    if w_match:
        w_power = int(w_match.group(1))
    # If there's a v or w without power, assign power as 1
    if 'v' in term and v_power == 0:
        v_power = 1
    if 'w' in term and w_power == 0:
        w_power = 1
    return v_power, w_power

def extract_coefficient(term):
    """Extract coefficient from a term."""
    # Checking for terms that start with 'v' or 'w' without any numeral in front
    if term.startswith('v') or term.startswith('w'):
        return 1
    match = re.search(r'(\d+)\*', term)
    if match:
        return int(match.group(1))
    return 1

  
def vw_powers(v_st,w_st,n):
    h1,h2,c=0,0,0

    if (v_st==2 or v_st==3) and w_st>0: #enojni heliks
        if w_st >3:
            h2= (w_st + 3) - 6
        else:
            h2 =0
        h1 = (w_st + 3) - h2
    else:
        h1=0
        h2=0
                
    c = (n+1) - (h1+h2)
        
    return (h1,h2,c)# potence (Y,X,v,w)


polys={7:"3*v**4*w**2 + 9*v**4*w + 19*v**4 + 2*v**3*w**3 + 6*v**3*w**2 + 12*v**3*w + 30*v**3 + v**2*w**5 + 2*v**2*w**4 + 3*v**2*w**3 + 4*v**2*w**2 + 5*v**2*w + 21*v**2 + 7*v + 1",
       17:"13*v**4*w**12 + 39*v**4*w**11 + 78*v**4*w**10 + 130*v**4*w**9 + 195*v**4*w**8 + 273*v**4*w**7 + 364*v**4*w**6 + 468*v**4*w**5 + 585*v**4*w**4 + 715*v**4*w**3 + 858*v**4*w**2 + 1014*v**4*w + 2184*v**4 + 2*v**3*w**13 + 6*v**3*w**12 + 12*v**3*w**11 + 20*v**3*w**10 + 30*v**3*w**9 + 42*v**3*w**8 + 56*v**3*w**7 + 72*v**3*w**6 + 90*v**3*w**5 + 110*v**3*w**4 + 132*v**3*w**3 + 156*v**3*w**2 + 182*v**3*w + 665*v**3 + v**2*w**15 + 2*v**2*w**14 + 3*v**2*w**13 + 4*v**2*w**12 + 5*v**2*w**11 + 6*v**2*w**10 + 7*v**2*w**9 + 8*v**2*w**8 + 9*v**2*w**7 + 10*v**2*w**6 + 11*v**2*w**5 + 12*v**2*w**4 + 13*v**2*w**3 + 14*v**2*w**2 + 15*v**2*w + 136*v**2 + 17*v + 1",
       22:"378*v**4*w**12 + 504*v**4*w**11 + 648*v**4*w**10 + 810*v**4*w**9 + 990*v**4*w**8 + 1188*v**4*w**7 + 1404*v**4*w**6 + 1638*v**4*w**5 + 1890*v**4*w**4 + 2160*v**4*w**3 + 2448*v**4*w**2 + 2754*v**4*w + 6954*v**4 + 2*v**3*w**18 + 6*v**3*w**17 + 12*v**3*w**16 + 20*v**3*w**15 + 30*v**3*w**14 + 42*v**3*w**13 + 56*v**3*w**12 + 72*v**3*w**11 + 90*v**3*w**10 + 110*v**3*w**9 + 132*v**3*w**8 + 156*v**3*w**7 + 182*v**3*w**6 + 210*v**3*w**5 + 240*v**3*w**4 + 272*v**3*w**3 + 306*v**3*w**2 + 342*v**3*w + 1520*v**3 + v**2*w**20 + 2*v**2*w**19 + 3*v**2*w**18 + 4*v**2*w**17 + 5*v**2*w**16 + 6*v**2*w**15 + 7*v**2*w**14 + 8*v**2*w**13 + 9*v**2*w**12 + 10*v**2*w**11 + 11*v**2*w**10 + 12*v**2*w**9 + 13*v**2*w**8 + 14*v**2*w**7 + 15*v**2*w**6 + 16*v**2*w**5 + 17*v**2*w**4 + 18*v**2*w**3 + 19*v**2*w**2 + 20*v**2*w + 231*v**2 + 22*v + 1",
       27:"23*v**4*w**22 + 69*v**4*w**21 + 138*v**4*w**20 + 230*v**4*w**19 + 345*v**4*w**18 + 483*v**4*w**17 + 644*v**4*w**16 + 828*v**4*w**15 + 1035*v**4*w**14 + 1265*v**4*w**13 + 1518*v**4*w**12 + 1794*v**4*w**11 + 2093*v**4*w**10 + 2415*v**4*w**9 + 2760*v**4*w**8 + 3128*v**4*w**7 + 3519*v**4*w**6 + 3933*v**4*w**5 + 4370*v**4*w**4 + 4830*v**4*w**3 + 5313*v**4*w**2 + 5819*v**4*w + 16974*v**4 + 2*v**3*w**23 + 6*v**3*w**22 + 12*v**3*w**21 + 20*v**3*w**20 + 30*v**3*w**19 + 42*v**3*w**18 + 56*v**3*w**17 + 72*v**3*w**16 + 90*v**3*w**15 + 110*v**3*w**14 + 132*v**3*w**13 + 156*v**3*w**12 + 182*v**3*w**11 + 210*v**3*w**10 + 240*v**3*w**9 + 272*v**3*w**8 + 306*v**3*w**7 + 342*v**3*w**6 + 380*v**3*w**5 + 420*v**3*w**4 + 462*v**3*w**3 + 506*v**3*w**2 + 552*v**3*w + 2900*v**3 + v**2*w**25 + 2*v**2*w**24 + 3*v**2*w**23 + 4*v**2*w**22 + 5*v**2*w**21 + 6*v**2*w**20 + 7*v**2*w**19 + 8*v**2*w**18 + 9*v**2*w**17 + 10*v**2*w**16 + 11*v**2*w**15 + 12*v**2*w**14 + 13*v**2*w**13 + 14*v**2*w**12 + 15*v**2*w**11 + 16*v**2*w**10 + 17*v**2*w**9 + 18*v**2*w**8 + 19*v**2*w**7 + 20*v**2*w**6 + 21*v**2*w**5 + 22*v**2*w**4 + 23*v**2*w**3 + 24*v**2*w**2 + 25*v**2*w + 351*v**2 + 27*v + 1",
       32:"28*v**4*w**27 + 84*v**4*w**26 + 168*v**4*w**25 + 280*v**4*w**24 + 420*v**4*w**23 + 588*v**4*w**22 + 784*v**4*w**21 + 1008*v**4*w**20 + 1260*v**4*w**19 + 1540*v**4*w**18 + 1848*v**4*w**17 + 2184*v**4*w**16 + 2548*v**4*w**15 + 2940*v**4*w**14 + 3360*v**4*w**13 + 3808*v**4*w**12 + 4284*v**4*w**11 + 4788*v**4*w**10 + 5320*v**4*w**9 + 5880*v**4*w**8 + 6468*v**4*w**7 + 7084*v**4*w**6 + 7728*v**4*w**5 + 8400*v**4*w**4 + 9100*v**4*w**3 + 9828*v**4*w**2 + 10584*v**4*w + 35119*v**4 + 2*v**3*w**28 + 6*v**3*w**27 + 12*v**3*w**26 + 20*v**3*w**25 + 30*v**3*w**24 + 42*v**3*w**23 + 56*v**3*w**22 + 72*v**3*w**21 + 90*v**3*w**20 + 110*v**3*w**19 + 132*v**3*w**18 + 156*v**3*w**17 + 182*v**3*w**16 + 210*v**3*w**15 + 240*v**3*w**14 + 272*v**3*w**13 + 306*v**3*w**12 + 342*v**3*w**11 + 380*v**3*w**10 + 420*v**3*w**9 + 462*v**3*w**8 + 506*v**3*w**7 + 552*v**3*w**6 + 600*v**3*w**5 + 650*v**3*w**4 + 702*v**3*w**3 + 756*v**3*w**2 + 812*v**3*w + 4930*v**3 + v**2*w**30 + 2*v**2*w**29 + 3*v**2*w**28 + 4*v**2*w**27 + 5*v**2*w**26 + 6*v**2*w**25 + 7*v**2*w**24 + 8*v**2*w**23 + 9*v**2*w**22 + 10*v**2*w**21 + 11*v**2*w**20 + 12*v**2*w**19 + 13*v**2*w**18 + 14*v**2*w**17 + 15*v**2*w**16 + 16*v**2*w**15 + 17*v**2*w**14 + 18*v**2*w**13 + 19*v**2*w**12 + 20*v**2*w**11 + 21*v**2*w**10 + 22*v**2*w**9 + 23*v**2*w**8 + 24*v**2*w**7 + 25*v**2*w**6 + 26*v**2*w**5 + 27*v**2*w**4 + 28*v**2*w**3 + 29*v**2*w**2 + 30*v**2*w + 496*v**2 + 32*v + 1"
       }


# DEFINES the size of VW matrix
v_pow_max,w_pow_max=4,30 #max powers of v, w for v**n x w**m matrix -SAME for all the polynomials
T0=273.15
R=1.987*10**(-3)

"""
SPECTROSCOPIC WEIGHTS matrix generator
"""
def spectro_matrices(v_Pow_max,w_Pow_max): #RETURNS matrices to calculate contribution of each P term
    
    matrix_h1 = np.zeros((v_pow_max+1 ,w_pow_max+1), dtype=int)
    matrix_h2 = np.zeros((v_pow_max+1 ,w_pow_max+1), dtype=int)
    matrix_c = np.zeros((v_pow_max+1 ,w_pow_max+1), dtype=int)
    
    for vv in range(v_Pow_max+1):
        for ww in range(w_Pow_max+1):
            h1,h2,c=vw_powers(vv,ww,100)
            matrix_h1[vv,ww]=h1
            matrix_h2[vv,ww]=h2
            matrix_c[vv,ww]=c
            
    return (matrix_h1,matrix_h2,matrix_c)


matrix_h1,matrix_h2,matrix_c = spectro_matrices(v_pow_max, w_pow_max)

#CONTRIBUTIONS of EACH TERM in POLYNOMIAL
M_H2=matrix_h2.transpose().flatten()
M_H1=matrix_h1.transpose().flatten()
M_C=(matrix_c.transpose().flatten()-100)

"""
POLYNOMIAL to MATRIX REPRESENTATION
"""
def polynomial_to_matrix(poly,v_Pow_max,w_Pow_max):
    """Convert polynomial string to matrix of coefficients."""
    terms = poly.split('+')
    max_v_power, max_w_power = v_Pow_max,w_Pow_max

    matrix = np.zeros((max_v_power + 1,max_w_power + 1), dtype=int)

    for term in terms:
        term=term.replace(" ","")
        v_power, w_power = extract_powers(term)
        coefficient = extract_coefficient(term)

        if v_power==0 and w_power==0:
            matrix[v_power,w_power] = int(term)
        else:
            matrix[v_power,w_power] = coefficient

    return matrix

def matrix_vw(v,w,v_pow_max,w_pow_max):

    """        w_j**k                 """
    w_powers = np.arange(w_pow_max + 1)
    w_pow = np.power(w[:, None], w_powers)
    
    """         v**i                  """
    v_powers = np.arange(v_pow_max + 1)
    v_pow = np.power(v, v_powers);   
       
    """     (v**i)*(w_j)**k           """
    wv = v_pow[:,None,None]*w_pow
    wv_reshaped = wv.transpose(1, 2, 0).reshape(len(w), -1)
    
    return wv_reshaped



N=17 #choose function - polynomial
temps=np.array([i for i in range(100)])

dG,dH,dCp,v=-0.17,-1.0,0.008,0.07
w = np.exp((-1/R) * (dG/T0 + dH * ((1/(temps + 273.15)) - (1/T0)) + dCp * (1-(T0/(temps + 273.15)) - np.log((temps + 273.15)/T0))))
H2,dH2,H1,dH1,C,dC=-40000,250,-15000,10,2000,-50

VW = matrix_vw(v,w,v_pow_max,w_pow_max)

"""Coefficient matrix"""
coef_matrix = polynomial_to_matrix(polys[N],v_pow_max,w_pow_max)

"""Coefficient matrix * WV - FINAL matrix polynomial"""
matrix_poly = coef_matrix.transpose().flatten()*VW
P_matrix = matrix_poly/np.sum(matrix_poly,axis=1)[:,None] # each term divided by total sum of polynomial- normalization -> PROBABILITIES

#FINAL MODEL FUNCTION - M_H2,M_H1,M_C charachteristic matrices
cd = P_matrix*( M_H2[None,:]*(H2 + dH2*temps)[:,None] +
        M_H1[None,:]*(H1 + dH1*temps)[:,None] +
        (M_C+N)[None,:] *(C  + dC *temps)[:,None])

CD=cd.sum(axis=1)/(N+1) #result

plt.plot(temps,CD)

The matrix_vw function will work as-is, you just need to replace np. with pt.

For the polynomial_to_matrix function, you need to replace this:

matrix[v_power, w_power] = ...

With this:

matrix = pt.set_subtensor(matrix[v_power, w_power], ...)

This is because pytensor doesn’t support direct element assignment.

The calculations after that should work fine just switching np. to pt.

Thanks, that seems to work fine. However, when I try to implement this matrix calculation procedure to β€œpm.Model” section of the script there are some issues with reshaping of the tensor. (ValueError: The output dimensions after reshape must be 0 or greater). β†’ for wv_reshaped = wv.transpose(1, 2, 0).reshape(w.shape[0],-1)

Even if I explicitly define the desired (correct) shape of the output matrix WV_reshapes, reshaping does not seems to work correctly.

...rest of the code...

def matrix_vw(v,w,v_pow_max,w_pow_max):

    """        w_j**k                 """
    w_powers = np.arange(w_pow_max + 1)
    w_pow = at.power(w[:, None], w_powers)
    
    """         v**i                  """
    v_powers = np.arange(v_pow_max + 1)
    v_pow = at.power(v, v_powers);   
       
    """     (v**i)*(w_j)**k           """
    wv = v_pow[:, None, None]*w_pow
    
    print("WV",wv.shape)
    print("WV_T",wv.transpose(1, 2, 0).shape)
    print("WV_reshaped",wv.transpose(1, 2, 0).reshape(101,155).shape)

    return wv.transpose(1, 2, 0).reshape(101,155)

...

with pm.Model() as model:
    #VARIABLES
    dg = pm.Normal('dg', mu=-0.20, sigma=0.2)
    dh = pm.Normal('dh', mu=-1.0, sigma=0.2)
    dcp = pm.Normal('dcp', mu=0, sigma=0.01)
    v = pm.Normal('v', mu=0.06, sigma=0.02)
    H2=pm.Normal('H2', mu=-40000, sigma=5000)
    dH2t=pm.Normal('dH2', mu=1000, sigma=300)
    C=pm.Normal('C', mu=2200, sigma=500)
    C_temp=pm.Normal('C_temp', mu=-50, sigma=50)
    k=pm.Normal('k', mu=3.8, sigma=1)
    
    #Determenistic variables
    w = pm.math.exp((-1 / R) * (dg / T0 + dh * ((1 / T_CD_K) - (1 / T0)) + dcp * (1 - (T0 / T_CD_K) - pm.math.log(T_CD_K/ T0))))
    H2_t=H2 + dH2t*T_CD
    H1_t=H2_t*(1-k/6)
    C_t=C+C_temp*T_CD
    
    
    VW=matrix_vw(v, w, v_pow_max, w_pow_max)
    coef_matrix = polynomial_to_matrix(polys[N],v_pow_max,w_pow_max)
    
    matrix_poly = coef_matrix.transpose().flatten()*VW

    P_matrix = matrix_poly/at.sum(matrix_poly)[:,None] # each term divided by total sum of polynomial- normalization -> PROBABILITIES
    
    cd =    P_matrix*( M_H2[None,:]*(H2 + dH2t*T_CD)[:,None] +
        M_H1[None,:]*(H1_t)[:,None] +
        (M_C+N)[None,:] *(C+C_temp*T_CD)[:,None])
    
    mu = cd.sum(axis=1)/(N+1) #final model
    sd = pm.HalfNormal('sd', sigma=300)
    
    obs = pm.Normal('obs', mu=mu, sigma=sd, observed=Y)

    #Perform MCMC
    trace = pm.sample(10000, tune=4000,nuts_sampler= "numpyro")
type or paste code here

Prints:

WV [ 5 101 31]
WV_T [101 31 5]
WV_reshaped [101]

However WV_reshaped should have dimensions of (101,155).