Hi! It seems that scan does not support subsampling although I have implemented it by figuring the init parameter of scan with subsample size. So HMCECS does not recognize subsampling. I have to use the sequential form of my model to make sure it worked properly in HMCECS. Is there any way to let subsampling work in scan? I really appreciate it!
Both scan primitive and HMCECS are complicated to implement. It seems to me that it’s tricky to make them compatible. Can you rewrite your model to use jax’s scan instead?
Sure! I can make the scan form of my model work (without subsampling, and on a small size of subsampled dataset). However I have a large dataset so I have to considering using HMCECS to scale my model to this dataset.
Which scan operator did you mention in the title? HMCECS does not work with numpyro’s scan but I think it will work with jax’s scan.
Sorry I misunderstood your last reply. But actually, when I was using for loop in my model, HMCECS still worked will and the sample speed is acceptable. I visualized the sampled result just now and found that it seemed not converged. I didn’t apply Taylor proxy with which I’m not familiar. I found that faster sampling seemed to mean that the sampling was about to converge. But in my model, the sampling rate is always slow, about 10s/it, until the sampling ends. I set warmup=12000 and sample=3000. Do you think more samples are needed for a complex model?
The scan operator I mentioned is numpyro.contrib.control_flow.scan
I see. If you rewrite your model/likelihood using jax.scan (with numpyro.factor statements), then I think HMCECS will work.
Re sampling behavior: it’s a research topic, I guess. It’s better to reach out to HMCECS authors for recommendations.