Help with reduce_sum - hierarchical logistic regression

I am trying to convert the following model code to a multithreaded format using reduce_sum

Original:

data {
  int<lower=0> N;//Number of observations
  int<lower=1> J;//Number of predictors with random slope
  int<lower=1> K;//Number of predictors with non-random slope
  int<lower=1> L;//Number of customers/groups
  int<lower=0,upper=1> y[N];//Binary response variable
  int<lower=1,upper=L> ll[N];//Number of observations in groups
  matrix[N,K] x1;
  matrix[N,J] x2;
}
parameters {
  vector[J] rbeta_mu; //mean of distribution of beta parameters
  vector<lower=0>[J] rbeta_sigma; //variance of distribution of beta parameters
  vector[J] beta_raw[L]; //group-specific parameters beta
  vector[K] beta;
}
transformed parameters {
  vector[J] rbeta[L];
  for (l in 1:L)
    for (j in 1:J)
      rbeta[l][j] = rbeta_mu[j] + rbeta_sigma[j] * beta_raw[l][j]; // coefficients on x
}
model {
  rbeta_mu ~ normal(0,100);
  rbeta_sigma ~ cauchy(0,10);
  beta~normal(0,10);
  for (l in 1:L)
    beta_raw[l] ~ normal(0,1);

  for(n in 1:N)
    y[n]~bernoulli_logit(x1[n] * beta + x2[n] * rbeta[ll[n]]);
}

Reduce_Sum attempt :

functions {
  real partial_sum(int start, 
                   int end,
                   int[] y_slice, 
                   matrix x1,
                   matrix x2,
                   vector[] beta,
                   matrix rbeta,
                   int J, 
                   int[] ll) {

    return bernoulli_logit_lpmf(y_slice[start:end] | x1[start:end,] * beta + (x2[start:end,:] .* rbeta[ll[start:end],:]) * rep_vector(1,J));
  }
}
data {
  int<lower=0> N;//Number of observations
  int<lower=1> J;//Number of predictors with random slope
  int<lower=1> K;//Number of predictors with non-random slope
  int<lower=1> L;//Number of customers/groups
  int<lower=0,upper=1> y[N];//Binary response variable
  int<lower=1,upper=L> ll[N];//Number of observations in groups
  matrix[N,K] x1;
  matrix[N,J] x2;
}
parameters {
  vector[J] rbeta_mu; //mean of distribution of beta parameters
  vector<lower=0>[J] rbeta_sigma; //variance of distribution of beta parameters
  vector[J] beta_raw[L]; //group-specific parameters beta
  vector[K] beta;
}
transformed parameters {
  vector[J] rbeta[L];
  for (l in 1:L)
    for (j in 1:J)
      rbeta[l][j] = rbeta_mu[j] + rbeta_sigma[j] * beta_raw[l][j]; // coefficients on x
}
model {
  rbeta_mu ~ normal(0,100);
  rbeta_sigma ~ cauchy(0,10);
  beta~normal(0,10);
  for (l in 1:L)
    beta_raw[l] ~ normal(0,1);
  target+=reduce_sum(partial_sum,1,y,x1,x2,beta,rbeta,J,ll);
}

I am getting the following error:

This is the statement giving the error:

    return bernoulli_logit_lpmf(y_slice[start:end] | x1[start:end,] * beta + (x2[start:end,:] .* rbeta[ll[start:end],:]) * rep_vector(1,J));

Any help on how to resolve it is greatly appreciated.

The error message seems a bit weird, but there is at least one error I am seeing:

x2[start:end,:] .* rbeta[ll[start:end],:]

The problem is that:

  • x2[start:end, :] is a matrix
  • rbeta[ll[start:end],:] is an array of vector

The product of these 2 types is not defined.
Its not really evident what is desired here, but I am guessing elementwise products of matrix rows and vectors? If so, that has to be done with a loop.

What if rbeta was a matrix :

functions {
  real partial_sum(int start, 
                   int end,
                   int[] y_slice, 
                   matrix x1,
                   matrix x2,
                   vector[] beta,
                   matrix rbeta,
                   int J, 
                   int[] ll) {

    return bernoulli_logit_lpmf(y_slice[start:end] | x1[start:end,] * beta + (x2[start:end,:] .* rbeta[ll[start:end],:]) * rep_vector(1,J));
  }
}
data {
  int<lower=0> N;//Number of observations
  int<lower=1> J;//Number of predictors with random slope
  int<lower=1> K;//Number of predictors with non-random slope
  int<lower=1> L;//Number of customers/groups
  int<lower=0,upper=1> y[N];//Binary response variable
  int<lower=1,upper=L> ll[N];//Number of observations in groups
  matrix[N,K] x1;
  matrix[N,J] x2;
}
parameters {
  vector[J] rbeta_mu; //mean of distribution of beta parameters
  vector<lower=0>[J] rbeta_sigma; //variance of distribution of beta parameters
  vector[J] beta_raw[L]; //group-specific parameters beta
  vector[K] beta;
}
transformed parameters {
  matrix[L,J] rbeta;
  for (l in 1:L)
    rbeta[l] = rbeta_mu + rbeta_sigma .* beta_raw[l]; // coefficients on x
}
model {
  rbeta_mu ~ normal(0,100);
  rbeta_sigma ~ cauchy(0,10);
  beta~normal(0,10);
  for (l in 1:L)
    beta_raw[l] ~ normal(0,1);
  target+=reduce_sum(partial_sum,1,y,x1,x2,beta,rbeta,J,ll);
}

@rok_cesnovar even upon converting rbeta to matrix it does not help.

@rok_cesnovar

Upon just adding some brackets the error now changes position :

Before:

After:

The only difference between the two is the paranthesis around x1[start:end,:] * beta nothing else and here this should be possible.

A Bernoulli logit is not a good candidate for reduce_sum is what I would expect… but you can try.

In any case, maybe you consider to rewrite your model using the R package brms. Then you can just turn on a switch with brm for threading. Have a look for the vignette on threading within brms.