Parallel non-looped non-linear functions are inconsistent

[wds15]

When turning on threading for non-looped non-linear models, the provided non-linear function will get a mix of partial and full data structures. The partial log lik function generated looks like this (see example code below):

  // compute partial sums of the log-likelihood
  real partial_log_lik_lpmf(int[] seq, int start, int end, vector Y, matrix X_ult, vector b_ult, matrix X_omega, vector b_omega, matrix X_theta, vector b_theta, int[] C_1, int[] C_2, int[] C_3, real sigma, int[] J_1, vector Z_1_ult_1, vector r_1_ult_1) {
    real ptarget = 0;
    int N = end - start + 1;
    // initialize linear predictor term
    vector[N] nlp_ult = X_ult[start:end] * b_ult;
    // initialize linear predictor term
    vector[N] nlp_omega = X_omega[start:end] * b_omega;
    // initialize linear predictor term
    vector[N] nlp_theta = X_theta[start:end] * b_theta;
    // initialize non-linear predictor term
    vector[N] mu;
    for (n in 1:N) {
      // add more terms to the linear predictor
      int nn = n + start - 1;
      nlp_ult[n] += r_1_ult_1[J_1[nn]] * Z_1_ult_1[nn];
    }
    // compute non-linear predictor values
    mu = growth_test(nlp_ult , C_1 , nlp_theta , nlp_omega , C_2 , C_3);
    ptarget += normal_lpdf(Y[start:end] | mu, sigma);
    return ptarget;
  }

The issue is that now the parameters have a size according to the partial sum block while the covariates span the entire data-set (and the non-linear function does not get start and end; this cannot figure out how things line up). I think a consistent solution would be this:

  // compute partial sums of the log-likelihood
  real partial_log_lik_lpmf(int[] seq, int start, int end, vector Y, matrix X_ult, vector b_ult, matrix X_omega, vector b_omega, matrix X_theta, vector b_theta, int[] C_1, int[] C_2, int[] C_3, real sigma, int[] J_1, vector Z_1_ult_1, vector r_1_ult_1) {
    real ptarget = 0;
    int N = end - start + 1;
    // initialize linear predictor term
    vector[N] nlp_ult = X_ult[start:end] * b_ult;
    // initialize linear predictor term
    vector[N] nlp_omega = X_omega[start:end] * b_omega;
    // initialize linear predictor term
    vector[N] nlp_theta = X_theta[start:end] * b_theta;
    // initialize non-linear predictor term
    vector[N] mu;
    for (n in 1:N) {
      // add more terms to the linear predictor
      int nn = n + start - 1;
      nlp_ult[n] += r_1_ult_1[J_1[nn]] * Z_1_ult_1[nn];
    }
    // compute non-linear predictor values
    mu = growth_test(nlp_ult , C_1[start:end], nlp_theta , nlp_omega , C_2[start:end] , C_3[start:end]); // <-- here the covariates should be subsetted to start:end to have the same length as the other variables
    ptarget += normal_lpdf(Y[start:end] | mu, sigma);
    return ptarget;
  }

Here is the full toy example:

library(brms)
library(cmdstanr)

loss_alt <- transform(loss, row=as.integer(1:nrow(loss)), nr=nrow(loss), test=as.integer(0))
  scode_growth <- "
    vector growth_test(vector ult, int[] dev, vector theta, vector omega, int[] row, int[] test) {
      int N = rows(ult);
      vector[N] mu;
      int rows_sorted = 1;

      for(i in 1:N) {
           mu[i] = ult[i] * (1 - exp(-(dev[i]/theta[i])^omega[i]));
      }
      return(mu);
    }
  "
  growth_model  <- stanvar(name="growth_test", scode=scode_growth, block="functions")

  fit_loss <- make_stancode(
    bf(cum ~ growth_test(ult, dev, theta, omega, row, test),
       ult ~ 1 + (1|AY), omega ~ 1, theta ~ 1,
       nl = TRUE, loop=FALSE),
    data = loss_alt, family = gaussian(),
    stanvars = growth_model,
    prior = c(prior(normal(5000, 1000), nlpar = "ult"),
              prior(normal(1, 2), nlpar = "omega"),
              prior(normal(45, 10), nlpar = "theta")),
    control = list(adapt_delta = 0.9),
    refresh = 0,
    chains = 1,
    backend = "cmdstanr",
    threads = threading(2)
  )

cat(fit_loss, file="brms_reduce_sum_example_bug.stan")

[paul-buerkner]

Good point. Let me try to fix this.

[paul-buerkner]

This issue should now be fixed. Could you try it out to confirm the new behavior is as intended?

[wds15]

Whow... that was fast! The toy example from above results in the code as I suggested above. All good.

[paul-buerkner]

Thanks for checking!