Hello,
I am currently working on a multinomial regression model with many covariates and a large dataset. I am interested in reducing runtime by using within-chain parallelization. I initially thought the multinomial was a good candidate to use reduce_sum(). However, I don’ t seem to find a way that works and I got to a point where the model compiles properly but the chains don’ t start and the error messages are more puzzling than helpful.
This is as far as I got, with a very simple example for clarity:
Stan code:
// multinomial.stan
functions{
real partial_sum(int[,] y_slice,
int start,
int end,
vector[] theta) {
real interm_sum = 0;
for (i in start:end){
interm_sum += multinomial_lpmf(y_slice[i] | theta[i]);
}
return interm_sum;
}
}
data {
int<lower=0> Ny;
int<lower=0> K;
int<lower=0> Y_obs[K,Ny];
int grainsize;
}
parameters {
simplex[Ny] theta[K];
}
model {
target += reduce_sum(partial_sum,
Y_obs,
grainsize,
theta);
}
R code:
data <- t(rmultinom(5,500,c(0.5,0.3,0.2)))
stan_data <- list(Ny = 3,
K = 5,
Y_obs = data,
grainsize = 1)
mod <- cmdstan_model("./multinomial.stan",
cpp_options = list(stan_threads = TRUE))
fit3 <- mod$sample(data = stan_data,
seed = 5446,
chains = 3,
parallel_chains = 3,
threads_per_chain = 2,
iter_warmup = 500,
iter_sampling = 500,
refresh = 100)
here is part of the error message:
Running MCMC with 3 parallel chains, with 2 thread(s) per chain...
Chain 1 Unrecoverable error evaluating the log probability at the initial value.
Chain 1 Exception: Exception: array[uni, ...] index: accessing element out of range. index 2 out of range; expecting index to be between 1 and 1 (in '/tmp/RtmpVuSgYx/model-da9c97f8797.stan', line 10, column 6 to column 60) (in '/tmp/RtmpVuSgYx/model-da9c97f8797.stan', line 30, column 2 to line 33, column 32)
Chain 1 Exception: Exception: array[uni, ...] index: accessing element out of range. index 2 out of range; expecting index to be between 1 and 1 (in '/tmp/RtmpVuSgYx/model-da9c97f8797.stan', line 10, column 6 to column 60) (in '/tmp/RtmpVuSgYx/model-da9c97f8797.stan', line 30, column 2 to line 33, column 32)
Chain 1 Exception: Exception: array[uni, ...] index: accessing element out of range. index 2 out of range; expecting index to be between 1 and 1 (in '/tmp/RtmpVuSgYx/model-da9c97f8797.stan', line 10, column 6 to column 60) (in '/tmp/RtmpVuSgYx/model-da9c97f8797.stan', line 30, column 2 to line 33, column 32)
The same error message is repeated for chains 2 and 3, I just truncated the output above.
I also tried hard-coding the multinomial log-likelihood equation instead of using the multinomial() function (this is what I do without parallelization because it allows for vectorisation), but I run into incompatibility issues with data and variable types.
Did anyone manage to find a workaround with multinomial_lpmf() or hard-coded log-likelihood components?
Thank you
Joane