inc_beta overpromotion

Description

While investigating the inc_beta function (see #1237) I saw an opportunity to reduce overpromotion without doing any math.

Consider inc_beta's reverse mode implementation:

class inc_beta_vvv_vari : public op_vvv_vari {
 public:
  ...
  void chain() {
    const double beta_ab = beta(avi_->val_, bvi_->val_);
    grad_reg_inc_beta(d_a, d_b, avi_->val_, bvi_->val_, cvi_->val_,
                      digamma(avi_->val_), digamma(bvi_->val_),
                      digamma(avi_->val_ + bvi_->val_), beta_ab);

    avi_->adj_ += adj_ * d_a;
    bvi_->adj_ += adj_ * d_b;
    cvi_->adj_ += adj_ * std::pow(1 - cvi_->val_, bvi_->val_ - 1)
                  * std::pow(cvi_->val_, avi_->val_ - 1) / beta_ab;
  }
};
inline var inc_beta(const var& a, const var& b, const var& c) {
  return var(new internal::inc_beta_vvv_vari(a.vi_, b.vi_, c.vi_));
}

That is the only reverse-mode overload. The other cases, like vvd, are implicitly promoted to var.

But note that the chain method factors into two updates, one for the first two arguments, and one for the third, so we can trivially produce vvd and ddv varis that do less work.

There are also two more functions, inc_beta_dda and inc_beta_ddb, that calculate the partials for the first argument and the second separately. (Whereas grad_reg_inc_beta calculates the partials for the first and the second arguments simultaneously.) It might be faster to implement vdd, dvd, vdv, and dvv forms that use inc_beta_dda and inc_beta_ddb, rather than grad_reg_inc_beta.

There is some opportunity for reducing overpromotion in the forward mode implementation too.

inc_beta has mix mode tests via expect_ad, so we can be sure our derivatives are still correct; but a performance comparison would probably be a good idea too.

Example

+class inc_beta_vdd_vari : public op_vdd_vari {
+ public:
+  inc_beta_vdd_vari(vari* avi, const double bd, const double cd)
+      : op_vdd_vari(inc_beta(avi->val_, bd, cd), avi, bd, cd) {}
+  void chain() {
+    double ad = avi_->val_;
+    avi_->adj_
+        += adj_ * inc_beta_dda(ad, bd_, cd_, digamma(ad), digamma(ad + bd_));
+  }
+};

Expected Output

Same but faster.

Current Version:

v3.0.0

Start at the inc_beta reverse-mode implementation and its existing mixed-mode expect_ad tests. Compare the vvv path with the proposed vdd, dvd, vdv, and dvv forms, including inc_beta_dda and inc_beta_ddb, then inspect forward-mode overloads for similar overpromotion. Done means specialized overloads preserve derivative tests and show a measurable performance improvement.