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().
