should rvar auto-drop dimensions on subsetting?

Hi!

rvars are a great thing and I have started using them a bit. I have run into issues when using it and there could be a nice solution but let me start with describing my setting. I wanted to transfer code from the transformed parameters block into R in order to eventually handle post-processing of my draws in R instead of Stan. Here are the respective Stan bits:

parameters {
  // the first 1:num_groups parameters are EX modelled while the
  // num_groups+1:2*num_groups are NEX
  vector[2] log_beta_raw[2*num_groups,num_comp];
  vector[num_inter] eta_raw[2*num_groups];

  // hierarchical priors
  vector[2] mu_log_beta[num_comp];
  // for differential discounting we allow the tau's to vary by
  // stratum (but not the means)
  vector<lower=0>[2] tau_log_beta_raw[num_strata,num_comp];
  cholesky_factor_corr[2] L_corr_log_beta[num_comp];

  vector[num_inter] mu_eta;
  vector<lower=0>[num_inter] tau_eta_raw[num_strata];
  cholesky_factor_corr[num_inter] L_corr_eta;
}
transformed parameters {
  vector[2] beta[2*num_groups,num_comp];
  vector[num_inter] eta[2*num_groups];
  vector<lower=0>[2] tau_log_beta[num_strata,num_comp];
  vector<lower=0>[num_inter] tau_eta[num_strata];

  // in the case of fixed tau's we fill them in here
  if (prior_tau_dist == 0) {
    tau_log_beta = prior_EX_tau_mean_comp;
    tau_eta = prior_EX_tau_mean_inter;
  } else {
    tau_log_beta = tau_log_beta_raw;
    tau_eta = tau_eta_raw;
  }

  // EX parameters which vary by stratum which is defined by the group
  for(g in 1:num_groups) {
    int s = group_stratum_cid[g];
    for(j in 1:num_comp) {
      beta[g,j] = mu_log_beta[j] +
        diag_pre_multiply(tau_log_beta[s,j], L_corr_log_beta[j]) * log_beta_raw[g,j];
    }
    if (num_inter > 0)
      eta[g] = mu_eta + diag_pre_multiply(tau_eta[s], L_corr_eta) * eta_raw[g];
  }

  // NEX parameters
  beta[num_groups+1:2*num_groups] = log_beta_raw[num_groups+1:2*num_groups];
  eta[num_groups+1:2*num_groups] = eta_raw[num_groups+1:2*num_groups];

  // exponentiate the slope parameter to force positivity
  for(g in 1:2*num_groups)
    for(j in 1:num_comp)
      beta[g,j,2] = exp(beta[g,j,2]);
}

I ended up with this respective rvar style code:

library(OncoBayes2)
library(posterior)

example_model("combo2")

post <- blrmfit$posterior

spost <- subset_draws(as_draws_rvars(post),
                      variable=c("log_beta_raw", "eta_raw", "mu_log_beta", "tau_log_beta_raw", "L_corr_log_beta", "mu_eta", "tau_eta_raw", "L_corr_eta"))

sdata <- blrmfit$standata

S_chain <- 1E3
num_chains <- 4
S <- S_chain * num_chains

beta <- rvar(array(0, dim=c(S, 2*sdata$num_groups,sdata$num_comp,2)), nchain=num_chains)
eta <- rvar(array(0, dim=c(S, 2*sdata$num_groups)), nchain=num_chains)
tau_log_beta <- rvar(array(0, dim=c(S, sdata$num_strata, sdata$num_comp)), nchain=num_chains)
tau_eta <- rvar(array(0, dim=c(S, sdata$num_strata)), nchain=num_chains)

## in the case of fixed tau's we fill them in here
if (sdata$prior_tau_dist == 0) {
    tau_log_beta = as_rvar(sdata$prior_EX_tau_mean_comp)
    tau_eta = as_rvar(sdata$prior_EX_tau_mean_inter)
} else {
    tau_log_beta = spost$tau_log_beta_raw
    tau_eta = spost$tau_eta_raw
}

## EX parameters which vary by stratum which is defined by the group
for(g in 1:sdata$num_groups) {
    s = sdata$group_stratum_cid[g];
    for(j in 1:sdata$num_comp) {
        beta[g,j,drop=TRUE] = spost$mu_log_beta[j,,drop=TRUE] + (rdo(diag(tau_log_beta[s,j,,drop=TRUE], sdata$num_comp, sdata$num_comp)) %**% spost$L_corr_log_beta[j,,drop=TRUE]) %**% spost$log_beta_raw[g,j,,drop=TRUE]
        ##diag_pre_multiply(tau_log_beta[s,j], L_corr_log_beta[j]) * log_beta_raw[g,j];
    }
    if (num_inter > 0)
        eta[g] = spost$mu_eta + (rdo(diag(tau_eta[s,drop=TRUE], sdata$num_inter, sdata$num_inter)) %**% spost$L_corr_eta) %**% spost$eta_raw[g]
    ##+ diag_pre_multiply(tau_eta[s], L_corr_eta) * eta_raw[g];
}


## NEX parameters
beta[sdata$num_groups+1:2*sdata$num_groups] = spost$log_beta_raw[sdata$num_groups+1:2*sdata$num_groups];
eta[sdata$num_groups+1:2*sdata$num_groups] = spost$eta_raw[sdata$num_groups+1:2*sdata$num_groups];

## exponentiate the slope parameter to force positivity
for(g in 1:2*sdata$num_groups)
    for(j in 1:sdata$num_comp)
        beta[g,j,2] = exp(beta[g,j,2]);

You see that the rvars code is not nice as I have to use many times drop. This is due to the fact that I nest matrices and vectors into arrays... but that structure gets cast into big arrays in rvar. However, when subsetting the leading indices (for the array stuff in Stan) and then using the rvar, the math won't work, since the dimensions do not line up - this is due to it's intended use is after subsetting things first to the array level.

Thus, I was wondering if we could have a as_draws_rvars flavour which creates nested lists from the samples? So all arrays get converted to nested lists and the vectors/matrices become rvars? With that I would not have to deal with drop in such an ugly way.

... or maybe there is a simpler way of doing this?

Start with the as_draws_rvars conversion and rvar subsetting behavior described in the issue, then reproduce the beta and eta calculations shown in the R example. Done would mean the requested nested-array handling or an agreed simpler approach avoids the repeated drop=TRUE calls while preserving those calculations.