StanHeaders compatibility for both Stan 2.26 and 2.21

Introduce `USE_STANC3` macro to branch changes specific to Stan 2.26
that cannot be supported by the older version of `stanc` transpiler.

Partially reverts stan-dev#2190.

Signed-off-by: Hamada S. Badr <[email protected]>

Author: hsbadr

stan/math/prim/err/elementwise_check.hpp

@@ -144,9 +144,14 @@ inline void elementwise_check(const F& is_good, const char* function,
  */
 template <typename F, typename T, typename... Indexings,
           require_eigen_t<T>* = nullptr,
+#ifdef USE_STANC3
           std::enable_if_t<(Eigen::internal::traits<T>::Flags
                             & Eigen::LinearAccessBit)
                            || T::IsVectorAtCompileTime>* = nullptr>
+#else
+          std::enable_if_t<static_cast<bool>(Eigen::internal::traits<T>::Flags&(
+              Eigen::LinearAccessBit | Eigen::DirectAccessBit))>* = nullptr>
+#endif
 inline void elementwise_check(const F& is_good, const char* function,
                               const char* name, const T& x, const char* must_be,
                               const Indexings&... indexings) {
@@ -194,13 +199,23 @@ inline void elementwise_check(const F& is_good, const char* function,
  * @throws `std::domain_error` if `is_good` returns `false` for the value
  * of any element in `x`
  */
+#ifdef USE_STANC3
 template <typename F, typename T, typename... Indexings,
           require_eigen_t<T>* = nullptr,
           std::enable_if_t<!(Eigen::internal::traits<T>::Flags
                              & Eigen::LinearAccessBit)
                            && !T::IsVectorAtCompileTime
                            && !(Eigen::internal::traits<T>::Flags
                                 & Eigen::RowMajorBit)>* = nullptr>
+#else
+template <
+    typename F, typename T, typename... Indexings,
+    require_eigen_t<T>* = nullptr,
+    std::enable_if_t<!(Eigen::internal::traits<T>::Flags
+                       & (Eigen::LinearAccessBit | Eigen::DirectAccessBit))
+                     && !(Eigen::internal::traits<T>::Flags
+                          & Eigen::RowMajorBit)>* = nullptr>
+#endif
 inline void elementwise_check(const F& is_good, const char* function,
                               const char* name, const T& x, const char* must_be,
                               const Indexings&... indexings) {
@@ -237,13 +252,23 @@ inline void elementwise_check(const F& is_good, const char* function,
  * @throws `std::domain_error` if `is_good` returns `false` for the value
  * of any element in `x`
  */
+#ifdef USE_STANC3
 template <typename F, typename T, typename... Indexings,
           require_eigen_t<T>* = nullptr,
           std::enable_if_t<
               !(Eigen::internal::traits<T>::Flags & Eigen::LinearAccessBit)
               && !T::IsVectorAtCompileTime
               && static_cast<bool>(Eigen::internal::traits<T>::Flags
                                    & Eigen::RowMajorBit)>* = nullptr>
+#else
+template <
+    typename F, typename T, typename... Indexings,
+    require_eigen_t<T>* = nullptr,
+    std::enable_if_t<!(Eigen::internal::traits<T>::Flags
+                       & (Eigen::LinearAccessBit | Eigen::DirectAccessBit))
+                     && static_cast<bool>(Eigen::internal::traits<T>::Flags
+                                          & Eigen::RowMajorBit)>* = nullptr>
+#endif
 inline void elementwise_check(const F& is_good, const char* function,
                               const char* name, const T& x, const char* must_be,
                               const Indexings&... indexings) {

stan/math/prim/fun/add.hpp

@@ -42,7 +42,11 @@ template <typename Mat1, typename Mat2,
           require_all_not_st_var<Mat1, Mat2>* = nullptr>
 inline auto add(const Mat1& m1, const Mat2& m2) {
   check_matching_dims("add", "m1", m1, "m2", m2);
+#ifdef USE_STANC3
   return m1 + m2;
+#else
+  return (m1 + m2).eval();
+#endif
 }
 
 /**
@@ -74,7 +78,11 @@ template <typename Scal, typename Mat, require_stan_scalar_t<Scal>* = nullptr,
           require_eigen_t<Mat>* = nullptr,
           require_all_not_st_var<Scal, Mat>* = nullptr>
 inline auto add(const Scal c, const Mat& m) {
+#ifdef USE_STANC3
   return (c + m.array()).matrix();
+#else
+  return (c + m.array()).matrix().eval();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/assign.hpp

@@ -41,8 +41,12 @@ inline void print_mat_size(std::ostream& o) {
  * @param x Left-hand side.
  * @param y Right-hand side.
  */
+#ifdef USE_STANC3
 template <typename T_lhs, typename T_rhs,
           require_all_stan_scalar_t<T_lhs, T_rhs>* = nullptr>
+#else
+template <typename T_lhs, typename T_rhs>
+#endif
 inline void assign(T_lhs& x, const T_rhs& y) {
   x = y;
 }
@@ -62,12 +66,59 @@ inline void assign(T_lhs& x, const T_rhs& y) {
  * @param y Right-hand side matrix.
  * @throw std::invalid_argument if sizes do not match.
  */
+#ifdef USE_STANC3
 template <typename T_lhs, typename T_rhs,
           require_all_eigen_t<T_lhs, T_rhs>* = nullptr>
 inline void assign(T_lhs&& x, const T_rhs& y) {
   check_matching_dims("assign", "left-hand-side", x, "right-hand-side", y);
   x = y.template cast<value_type_t<T_lhs>>();
 }
+#else
+template <typename T_lhs, typename T_rhs, int R, int C>
+inline void assign(Eigen::Matrix<T_lhs, R, C>& x,
+                   const Eigen::Matrix<T_rhs, R, C>& y) {
+   check_matching_dims("assign", "left-hand-side", x, "right-hand-side", y);
+  for (int i = 0; i < x.size(); ++i) {
+    assign(x(i), y(i));
+  }
+}
+
+/**
+ * Copy the right-hand side's value to the left-hand side
+ * variable.
+ *
+ * <p>The <code>assign()</code> function is overloaded.  This
+ * instance will be called for arguments that are both
+ * <code>Eigen::Matrix</code> types and whose shapes match.  The
+ * shape of the right-hand side matrix is specified in the row and
+ * column shape template parameters.
+ *
+ * <p>The left-hand side is intentionally not a reference, because
+ * that won't match generally enough;  instead, a non-reference is
+ * used, which still holds onto a reference to the contained
+ * matrix and thus still updates the appropriate values.
+ *
+ * @tparam T_lhs Type of matrix block elements.
+ * @tparam T Type of right-hand side matrix elements.
+ * @tparam R Row shape for right-hand side matrix.
+ * @tparam C Column shape for right-hand side matrix.
+ * @param x Left-hand side block view of matrix.
+ * @param y Right-hand side matrix.
+ * @throw std::invalid_argument if sizes do not match.
+ */
+template <typename T_lhs, typename T, int R, int C>
+inline void assign(Eigen::Block<T_lhs> x, const Eigen::Matrix<T, R, C>& y) {
+  check_size_match("assign", "left-hand side rows", x.rows(),
+                   "right-hand side rows", y.rows());
+  check_size_match("assign", "left-hand side cols", x.cols(),
+                   "right-hand side cols", y.cols());
+  for (int n = 0; n < y.cols(); ++n) {
+    for (int m = 0; m < y.rows(); ++m) {
+      assign(x(m, n), y(m, n));
+    }
+  }
+}
+#endif
 
 /**
  * Copy the right-hand side's value to the left-hand side

stan/math/prim/fun/col.hpp

@@ -23,7 +23,11 @@ namespace math {
 template <typename T, typename = require_eigen_t<T>>
 inline auto col(const T& m, size_t j) {
   check_column_index("col", "j", m, j);
+#ifdef USE_STANC3
   return m.col(j - 1);
+#else
+  return m.col(j - 1).eval();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/corr_matrix_free.hpp

@@ -39,17 +39,29 @@ Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, 1> corr_matrix_free(const T& y) {
   check_nonzero_size("corr_matrix_free", "y", y);
 
   Eigen::Index k = y.rows();
+#ifdef USE_STANC3
   Eigen::Index k_choose_2 = (k * (k - 1)) / 2;
+#else
+  Array<value_type_t<T>, Dynamic, 1> x(k_choose_2);
+#endif
   Eigen::Matrix<value_type_t<T>, Dynamic, 1> x(k_choose_2);
   Array<value_type_t<T>, Dynamic, 1> sds(k);
+#ifdef USE_STANC3
   bool successful = factor_cov_matrix(y, x.array(), sds);
+#else
+  bool successful = factor_cov_matrix(y, x, sds);
+#endif
   if (!successful) {
     throw_domain_error("corr_matrix_free", "factor_cov_matrix failed on y", y,
                        "");
   }
   check_bounded("corr_matrix_free", "log(sd)", sds, -CONSTRAINT_TOLERANCE,
                 CONSTRAINT_TOLERANCE);
+#ifdef USE_STANC3
   return x;
+#else
+  return x.matrix();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/diag_post_multiply.hpp

@@ -24,7 +24,11 @@ template <typename T1, typename T2, require_eigen_t<T1>* = nullptr,
 auto diag_post_multiply(const T1& m1, const T2& m2) {
   check_size_match("diag_post_multiply", "m2.size()", m2.size(), "m1.cols()",
                    m1.cols());
+#ifdef USE_STANC3
   return m1 * m2.asDiagonal();
+#else
+  return (m1 * m2.asDiagonal()).eval();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/diag_pre_multiply.hpp

@@ -24,7 +24,11 @@ template <typename T1, typename T2, require_eigen_vector_t<T1>* = nullptr,
 auto diag_pre_multiply(const T1& m1, const T2& m2) {
   check_size_match("diag_pre_multiply", "m1.size()", m1.size(), "m2.rows()",
                    m2.rows());
+#ifdef USE_STANC3
   return m1.asDiagonal() * m2;
+#else
+  return (m1.asDiagonal() * m2).eval();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/diagonal.hpp

@@ -17,7 +17,11 @@ namespace math {
  */
 template <typename T, typename = require_eigen_t<T>>
 inline auto diagonal(const T& m) {
+#ifdef USE_STANC3
   return m.diagonal();
+#else
+  return m.diagonal().eval();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/divide.hpp

@@ -19,7 +19,11 @@ namespace math {
  * @return Scalar divided by the scalar.
  */
 template <typename Scal1, typename Scal2,
+#ifdef USE_STANC3
           require_all_stan_scalar_t<Scal1, Scal2>* = nullptr>
+#else
+          typename = require_all_stan_scalar_t<Scal1, Scal2>>
+#endif
 inline return_type_t<Scal1, Scal2> divide(const Scal1& x, const Scal2& y) {
   return x / y;
 }
@@ -41,10 +45,19 @@ inline int divide(int x, int y) {
  * @return matrix divided by the scalar
  */
 template <typename Mat, typename Scal, typename = require_eigen_t<Mat>,
+#ifdef USE_STANC3
           require_stan_scalar_t<Scal>* = nullptr,
           require_all_not_var_t<scalar_type_t<Mat>, Scal>* = nullptr>
+#else
+          typename = require_stan_scalar_t<Scal>,
+          typename = require_all_not_var_t<scalar_type_t<Mat>, Scal>>
+#endif
 inline auto divide(const Mat& m, Scal c) {
+#ifdef USE_STANC3
   return m / c;
+#else
+  return (m / c).eval();
+#endif
 }
 
 }  // namespace math

stan/math/prim/fun/elt_divide.hpp

@@ -24,7 +24,11 @@ template <typename Mat1, typename Mat2,
           require_all_not_st_var<Mat1, Mat2>* = nullptr>
 auto elt_divide(const Mat1& m1, const Mat2& m2) {
   check_matching_dims("elt_divide", "m1", m1, "m2", m2);
+#ifdef USE_STANC3
   return (m1.array() / m2.array()).matrix();
+#else
+  return (m1.array() / m2.array()).matrix().eval();
+#endif
 }
 
 /**
@@ -38,11 +42,19 @@ auto elt_divide(const Mat1& m1, const Mat2& m2) {
  * @param s scalar
  * @return Elementwise division of a scalar by matrix.
  */
+#ifdef USE_STANC3
 template <typename Mat, typename Scal, require_matrix_t<Mat>* = nullptr,
           require_stan_scalar_t<Scal>* = nullptr>
 auto elt_divide(const Mat& m, Scal s) {
   return divide(m, s);
 }
+#else
+template <typename Scal, typename Mat, typename = require_stan_scalar_t<Scal>,
+          typename = require_eigen_t<Mat>>
+auto elt_divide(Scal s, const Mat& m) {
+  return (s / m.array()).matrix().eval();
+}
+#endif
 
 /**
  * Return the elementwise division of the specified scalar

stan/math/prim/fun/elt_multiply.hpp

@@ -24,7 +24,11 @@ template <typename Mat1, typename Mat2,
           require_all_not_st_var<Mat1, Mat2>* = nullptr>
 auto elt_multiply(const Mat1& m1, const Mat2& m2) {
   check_matching_dims("elt_multiply", "m1", m1, "m2", m2);
+#ifdef USE_STANC3
   return m1.cwiseProduct(m2);
+#else
+  return m1.cwiseProduct(m2).eval();
+#endif
 }
 
 /**

stan/math/prim/fun/head.hpp

@@ -23,7 +23,11 @@ inline auto head(const T& v, size_t n) {
   if (n != 0) {
     check_vector_index("head", "n", v, n);
   }
+#ifdef USE_STANC3
   return v.head(n);
+#else
+  return v.head(n).eval();
+#endif
 }
 
 /**

stan/math/prim/fun/lb_constrain.hpp

@@ -55,8 +55,13 @@ inline return_type_t<T, L> lb_constrain(const T& x, const L& lb) {
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
+#ifdef USE_STANC3
 template <typename T, typename L, typename S>
 inline return_type_t<T, L> lb_constrain(const T& x, const L& lb, S& lp) {
+#else
+template <typename T, typename L>
+inline return_type_t<T, L> lb_constrain(const T& x, const L& lb, T& lp) {
+#endif
   using std::exp;
   if (lb == NEGATIVE_INFTY) {
     return identity_constrain(x, lp);

stan/math/prim/fun/log_softmax.hpp

@@ -5,7 +5,9 @@
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
+#ifdef USE_STANC3
 #include <stan/math/prim/fun/to_ref.hpp>
+#endif
 #include <stan/math/prim/functor/apply_vector_unary.hpp>
 
 namespace stan {
@@ -39,6 +41,7 @@ namespace math {
  * @param[in] x vector to transform
  * @return log unit simplex result of the softmax transform of the vector.
  */
+#ifdef USE_STANC3
 template <typename Container, require_st_arithmetic<Container>* = nullptr>
 inline auto log_softmax(const Container& x) {
   check_nonzero_size("log_softmax", "v", x);
@@ -49,7 +52,21 @@ inline auto log_softmax(const Container& x) {
       },
       to_ref(x));
 }
-
+#else
+/**
+ * Note: The return must be evaluated otherwise the Ref object falls out
+ * of scope
+*/
+template <typename Container, require_st_arithmetic<Container>* = nullptr>
+inline auto log_softmax(const Container& x) {
+  check_nonzero_size("log_softmax", "v", x);
+  return apply_vector_unary<Container>::apply(x, [](const auto& v) {
+    const Eigen::Ref<const plain_type_t<decltype(v)>>& v_ref = v;
+    check_nonzero_size("log_softmax", "v", v_ref);
+    return (v_ref.array() - log_sum_exp(v_ref)).eval();
+  });
+}
+#endif
 }  // namespace math
 }  // namespace stan
 #endif