Unable to parameterize lag parameter

Hi @ricardoV94 Thanks for this. I did run a simple experiment to compare the distributions of random variables with and without eval() and I would love your thoughts on this. The distributions look very similar .The code I used is as follows:

import pymc as pm
import pymc.sampling.jax as pmjax
import numpy as np

x = [i for i in range(200)]
x_np = np.array(x)
with pm.Model() as m:
    
    z = pm.Normal('z',180)
    u = np.roll(x_np,int(z.eval()))
with m:

    idata = pmjax.sample_numpyro_nuts(1000,chains = 4)

#### Without eval()

with pm.Model() as m2:
    
    z = pm.Normal('z',180)

with m2:

    idatax = pmjax.sample_numpyro_nuts(1000,chains = 4)

import matplotlib.pyplot as plt


plt.hist(np.mean(idata.posterior['z'],axis=0),label='With eval()',color='red')
plt.hist(np.mean(idatax.posterior['z'],axis=0),label = 'Without eval()',color='blue')
plt.legend()
plt.show()

The distributions of random variable ā€˜z’ with and without eval().