Combine `HMCECS` with `HMM_ENUM`

Hi! I have already implemented my hmm-like model by following the steps in HMM_ENUM with NUTS, But the number of i.i.d samples in my dataset is so large that I can only run my model on a very small size dataset from which I can’t get any comprehensive results. Now I’m trying to use PositionalSharding in jax to utilize all of my GPUs. Meanwhile, I notice that HMCECS maybe a better choice because it supports subsampling. However, after reading the tutorial of HMCECS. I found there is nothing about how to combine HMCECS with HMM. So I’m wondering if I can do this combination. Hope for your reply, really appreciate.