update neg_binomial_2_lpmf to be more vectorized friendly

Author: SteveBronder

stan/math/prim/fun/select.hpp

@@ -112,7 +112,7 @@ template <typename T_true, typename T_false,
                                               plain_type_t<T_false>>,
           require_stan_scalar_t<T_true>* = nullptr,
           require_container_t<T_false>* = nullptr>
-inline ReturnT select(const bool c, const T_true y_true, T_false&& y_false) {
+inline ReturnT select(const bool c, const T_true& y_true, T_false&& y_false) {
   if (c) {
     return apply_scalar_binary(
         [](auto&& y_true_inner, auto&& y_false_inner) { return y_true_inner; },
@@ -140,13 +140,13 @@ inline ReturnT select(const bool c, const T_true y_true, T_false&& y_false) {
 template <typename T_bool, typename T_true, typename T_false,
           require_eigen_array_vt<std::is_integral, T_bool>* = nullptr,
           require_all_stan_scalar_t<T_true, T_false>* = nullptr>
-inline auto select(const T_bool c, const T_true y_true, const T_false y_false) {
+inline auto select(T_bool&& c, const T_true& y_true, const T_false& y_false) {
   using ret_t = return_type_t<T_true, T_false>;
-  return c
-      .unaryExpr([y_true, y_false](bool cond) {
+  return make_holder([y_true, y_false](auto&& c_) {
+      return std::forward<decltype(c_)>(c_).unaryExpr([y_true, y_false](bool cond) {
         return cond ? ret_t(y_true) : ret_t(y_false);
-      })
-      .eval();
+      });
+  }, std::forward<T_bool>(c));
 }
 
 /**
@@ -164,13 +164,23 @@ inline auto select(const T_bool c, const T_true y_true, const T_false y_false) {
 template <typename T_bool, typename T_true, typename T_false,
           require_eigen_array_t<T_bool>* = nullptr,
           require_any_eigen_array_t<T_true, T_false>* = nullptr>
-inline auto select(const T_bool c, const T_true y_true, const T_false y_false) {
+inline auto select(T_bool&& c, T_true&& y_true, T_false&& y_false) {
   check_consistent_sizes("select", "boolean", c, "left hand side", y_true,
                          "right hand side", y_false);
   using ret_t = return_type_t<T_true, T_false>;
-  return c.select(y_true, y_false).template cast<ret_t>().eval();
+  if constexpr (!std::is_same_v<std::decay_t<T_true>, std::decay_t<T_false>>) {
+    return make_holder([](auto&& c_, auto&& y_true_, auto&& y_false_) {
+      return std::forward<decltype(c_)>(c_).select(
+          std::forward<decltype(y_true_)>(y_true_),
+          std::forward<decltype(y_false_)>(y_false_));
+    }, std::forward<T_bool>(c), std::forward<T_true>(y_true), std::forward<T_false>(y_false));
+  } else {
+    return make_holder([](auto&& c_, auto&& y_true_, auto&& y_false_) {
+      return std::forward<decltype(c_)>(c_).select(std::forward<decltype(y_true_)>(y_true_),
+                                    std::forward<decltype(y_false_)>(y_false_)).template cast<ret_t>();
+    }, std::forward<T_bool>(c), std::forward<T_true>(y_true), std::forward<T_false>(y_false));
+}
 }
-
 }  // namespace math
 }  // namespace stan
 

stan/math/prim/prob/neg_binomial_2_lpmf.hpp

@@ -11,6 +11,7 @@
 #include <stan/math/prim/fun/scalar_seq_view.hpp>
 #include <stan/math/prim/fun/size.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/select.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
 #include <stan/math/prim/functor/partials_propagator.hpp>
 #include <cmath>
@@ -51,67 +52,40 @@ inline return_type_t<T_location, T_precision> neg_binomial_2_lpmf(
   T_partials_return logp(0.0);
   auto ops_partials = make_partials_propagator(mu_ref, phi_ref);
 
-  scalar_seq_view<T_n_ref> n_vec(n_ref);
-  scalar_seq_view<T_mu_ref> mu_vec(mu_ref);
-  scalar_seq_view<T_phi_ref> phi_vec(phi_ref);
-  size_t size_mu = stan::math::size(mu);
-  size_t size_phi = stan::math::size(phi);
-  size_t size_mu_phi = max_size(mu, phi);
-  size_t size_n_phi = max_size(n, phi);
-  size_t size_all = max_size(n, mu, phi);
-
-  VectorBuilder<true, T_partials_return, T_location> mu_val(size_mu);
-  for (size_t i = 0; i < size_mu; ++i) {
-    mu_val[i] = mu_vec.val(i);
-  }
-
-  VectorBuilder<true, T_partials_return, T_precision> phi_val(size_phi);
-  VectorBuilder<true, T_partials_return, T_precision> log_phi(size_phi);
-  for (size_t i = 0; i < size_phi; ++i) {
-    phi_val[i] = phi_vec.val(i);
-    log_phi[i] = log(phi_val[i]);
-  }
-
-  VectorBuilder<true, T_partials_return, T_location, T_precision> mu_plus_phi(
-      size_mu_phi);
-  VectorBuilder<true, T_partials_return, T_location, T_precision>
-      log_mu_plus_phi(size_mu_phi);
-  for (size_t i = 0; i < size_mu_phi; ++i) {
-    mu_plus_phi[i] = mu_val[i] + phi_val[i];
-    log_mu_plus_phi[i] = log(mu_plus_phi[i]);
+  auto n_vec = as_array_or_scalar(n_ref);
+  auto mu_vec = as_array_or_scalar(mu_ref);
+  auto phi_vec = as_array_or_scalar(phi_ref);
+  decltype(auto) mu_val = value_of(mu_vec);
+  decltype(auto) phi_val = value_of(phi_vec);
+  auto log_phi = log(phi_val);
+  auto mu_plus_phi = mu_val + phi_val;
+  auto log_mu_plus_phi = log(mu_plus_phi);
+  auto n_plus_phi = value_of(n_vec) + phi_val;
+  constexpr bool include_precision = include_summand<propto, T_precision>::value;
+  constexpr bool include_location = include_summand<propto, T_location>::value;
+  auto logp_calc = [&]() {
+    return -phi_val * (log1p(mu_val / phi_val))
+          - value_of(n_vec) * log_mu_plus_phi;
+  };
+  if constexpr (include_precision || include_location) {
+    if constexpr (include_precision && include_location) {
+      logp += sum(binomial_coefficient_log(n_plus_phi - 1, n_vec) + multiply_log(n_vec, mu_val) + logp_calc());
+    } else if constexpr (include_precision) {
+      logp += sum(binomial_coefficient_log(n_plus_phi - 1, n_vec) + logp_calc());
+    } else if constexpr (include_location) {
+      logp += sum(multiply_log(n_vec, mu_val) + logp_calc());
+    }
   }
-
-  VectorBuilder<true, T_partials_return, T_n, T_precision> n_plus_phi(
-      size_n_phi);
-  for (size_t i = 0; i < size_n_phi; ++i) {
-    n_plus_phi[i] = n_vec[i] + phi_val[i];
+  if constexpr (is_autodiff_v<T_location>) {
+    partials<0>(ops_partials) = n_vec / mu_val - (n_vec + phi_val) / mu_plus_phi;
   }
-
-  for (size_t i = 0; i < size_all; i++) {
-    if constexpr (include_summand<propto, T_precision>::value) {
-      logp += binomial_coefficient_log(n_plus_phi[i] - 1, n_vec[i]);
-    }
-    if constexpr (include_summand<propto, T_location>::value) {
-      logp += multiply_log(n_vec[i], mu_val[i]);
-    }
-    logp += -phi_val[i] * (log1p(mu_val[i] / phi_val[i]))
-            - n_vec[i] * log_mu_plus_phi[i];
-
-    if constexpr (is_autodiff_v<T_location>) {
-      partials<0>(ops_partials)[i]
-          += n_vec[i] / mu_val[i] - (n_vec[i] + phi_val[i]) / mu_plus_phi[i];
-    }
-    if constexpr (is_autodiff_v<T_precision>) {
-      T_partials_return log_term;
-      if (mu_val[i] < phi_val[i]) {
-        log_term = log1p(-mu_val[i] / mu_plus_phi[i]);
-      } else {
-        log_term = log_phi[i] - log_mu_plus_phi[i];
-      }
-      partials<1>(ops_partials)[i] += (mu_val[i] - n_vec[i]) / mu_plus_phi[i]
-                                      + log_term - digamma(phi_val[i])
-                                      + digamma(n_plus_phi[i]);
-    }
+  if constexpr (is_autodiff_v<T_precision>) {
+    auto log_term
+        = select(mu_val < phi_val, log1p(-mu_val / mu_plus_phi),
+                 log_phi - log_mu_plus_phi);
+    partials<1>(ops_partials) = (mu_val - value_of(n_vec)) / mu_plus_phi
+                                    + log_term - digamma(phi_val)
+                                    + digamma(n_plus_phi);
   }
   return ops_partials.build(logp);
 }