templated propto & vectorized distributions
Bob Carpenter
Alexandrescu's). It'll be easy to write a single
templated propto_log version for each distribution
by leaning on the optimizing compiler and a little
template metaprogramming.
The trick is to write a templated struct and override it
statically to determine if a type is an auto-dif variable.
The code's really simple:
template <typename T>
struct is_var {
enum { Yes = false; }
}
template <>
struct is_var<stan::agrad::var> {
enum { Yes = true; }
}
With these two, is_var<T>::Yes is true if
T is stan::agrad::var and false otherwise.
Given that it's resolved at compile time, if-then
tests involving these struct's enums will get compiled
away.
For instance, instead of using eight template override functions,
which is hard to maintain and doc, we just have:
template <typename T_out, typename T_success, typename T_failure>
inline typename promote_args<T_out,T_success,T_failure>::type
beta_propto_log(T_out theta, T_success alpha, T_failure beta) {
if (is_var<T_success>::Yes || is_var<T_fail>::Yes) {
return -lbeta(alpha,beta)
+ (alpha - 1) * log(theta)
+ (beta - 1) * log(1 - theta);
} else if (is_var<T_out>) {
return (alpha - 1) * log(theta)
+ (beta - 1) * log(1 - theta);
} else {
return 0;
}
}
We should also vectorize the outcomes:
template <typename T_out, typename T_success, typename T_failure>
inline typename promote_args<T_out,T_success,T_failure>::type
beta_propto_log(vector<T_out> theta, T_success alpha, T_failure beta) {
if (is_var<T_success>::Yes && is_var<T_fail>::Yes) {
var lp = -lbeta(alpha,beta) * theta.size();
var alpha_m1 = alpha - 1;
var beta_m1 = beta - 1;
for (unsigned int i = 0; i < theta.size(); ++i)
lp += alpha_m1 * log(theta[i])
+ beta_m1 * log(1 - theta[i]);
return lp;
} else if ...
...
}
Then we only call the log beta function once, saving
theta.size() - 1 calls to the beta function and its
gradient, each of which would have made several log gamma
and digamma evals.
In the BUGS-like language, you'd replace
for (n in 1:N)
theta[n] ~ beta(alpha,beta);
with the vectorized form:
theta ~ beta(alpha,beta);
The result is equivalent, but the latter's more efficient.
The declaration of theta in both cases should be:
double(0,1) theta[N];
- Bob
