Hi! I am trying to implement a hierarchical change-point model in stan. At the heart of the stan implementation is the linear-time example given in the Stan reference (https://mc-stan.org/docs/2_18/stan-users-guide/change-point-section.html)

I have a collection of N time-series of length T which are divided into two classes. I have a list of ints called `classes`

which allows me to label each of the N time-series as an element of one of the two classes. The two classes have different breakpoints but they have the same early and late rates. My `transformed parameters`

block:

```
matrix[T, k] lp;
matrix[T + 1, k] lp_e;
matrix[T + 1, k] lp_l;
matrix[T + 1, k] prior_s;
for (i in 1:k)
{
lp_e[1, i] = 0;
lp_l[1, i] = 0;
prior_s[1, i] = normal_lpdf(1 | mu, sigma);
}
for (t in 1:T) {
for (i in 1:k) {
prior_s[t + 1, i] = normal_lpdf(t | mu, sigma);
lp_e[t+1, i] = lp_e[t, i];
lp_l[t+1, i] = lp_l[t, i];
}
for (n in 1:N) {
lp_e[t+1, classes[n]] += normal_lpdf(D[n, t] | e, sigma_data);
lp_l[t+1, classes[n]] += normal_lpdf(D[n, t] | l, sigma_data);
}
}
for (i in 1:k) {
lp[:,i] = rep_vector(lp_l[T + 1, i], T )
+ head(lp_e[:, i], T) - head(lp_l[:, i], T) + prior_s[1:T, i];
}
```

I suspect the 1…N loop is slowing things down and I’d like to vectorise it. I will eventually have many more than two classes and N will be in the tens of thousands (at least). Any ideas on how to write an autodiff-friendly implementation? Many thanks!