Composable transforms in Stan

Developers
8

I’m not sure if this was brought up independently from the math meeting but we also spoke about composable constraint transforms. However, the idea was to expand the language with functions that would allow one to mix and match whichever transforms one wants, in the order given. My proposal is related to yours but can be considered independent. There would just be another way to accomplish the same task.

I’m going to work with @stevebronder to get a clear design doc put together. The idea though is that these functions would be callable in the transformed parameters or model blocks. The user would be able to tell a functor the constraint functions to apply in the order given and if a jacobian adjustment should be made.

It might be accomplished with something like

parameters {
 real rate_raw;
}
transformed parameters {
  real rate = constraint_transform(rate_raw, offset(mu), multiplier(sigma), lower(0), target);
}

Though this is way more verbose it allows opening up the constraint transforms to be used outside of the parameters block. These can then be mixed and matched as a user wants. For example, if a user wants a cholesky factor of correlation matrices where the corresponding entries of the correlation matrix are all constrained to be positive (see here for how to accomplish this today in Stan) the proposal would look something like (nb this is just an example of one way we could accomplish this, if it looks bad to you, please comment on the design doc to make it better),

parameters {
 vector<lower=0>[N * (N - 1) / 2] corr_raw;
}
transformed parameters {
 matrix[N, N] L = constraint_transform(corr_raw, corr(), cholesky_corr_constrain(), target);
}

Though one could also do

parameters {
 vector[N * (N - 1) / 2] corr_raw;
}
transformed parameters {
 matrix[N, N] L = constraint_transform(corr_raw, lower(0), corr(), cholesky_corr_constrain(), target);
}

We can see how the current constraint transform in the parameters block is utilized with this proposal. Using the lower bound declaration of the corr_raw vector ensures that only positive elements are given to corr_constrain. That transform is tanh(x) with the appropriate jacobian adjustment. The final cholesky_corr_constrain does what is described in 10.12 Cholesky factors of correlation matrices | Stan Reference Manual.

1 Like