Memory retention case study: behavior of generated quantities block and extract function

I am reading through “Bayesian cognitive modeling” and trying out the case studies using stan. In the case study on memory retention model, mem-ret_1.stan (1.1 KB) ,

The original github link for the stan code is https://github.com/stan-dev/example-models/blob/master/Bayesian_Cognitive_Modeling/CaseStudies/MemoryRetention/Retention_1_Stan.R.

The generated quantities block is as follows:
generated quantities {
int<lower=0, upper=n> predk[ns, nt];

// Predicted data
for (i in 1:ns) {
for (j in 1:nt) {
predk[i, j] = binomial_rng(n, theta[i,j]);
}
}
}

The prediction matrix is a 4 x 10 matrix, ns =4 and nt = 10. The “theta” parameter vector is also a 4x10 matrix. However when I extract the predk matrix after training the model using:

# generating a list of lists
inits <- list(alpha=0.5, beta=0.1)
inits_ll <- lapply(1:4, function(x) inits)

m <- stan(file = "mem-ret_1.stan",
          data = data,
          control = list(adapt_delta = 1 - 1e-3),
          cores = 4,
          chains = 4,
          init = init_ll,
          pars = c("alpha", "beta", "predk"))

rstan::extract(m)$predk[,,1]

The resulting predk matrix has 4 columns and 4000 rows. I was expecting to see a 4 rows with 10 columns. For the parameters alpha and beta I see 4000 values, which I understand. Can someone pls help me understand what I am missing?

The generated quantities block runs for each iteration, and the extract function returns all the those samples.

1 Like

thanks @emiruz, just noticed that the dimension of the output is, 4000x4x10. I was checking extract(m)$predk[,1].

dim(rstan::extract(m)$predk[,,])
[1] 4000   4    10