Broaden the template bound on the lp argument to constraints

Author: WardBrian

stan/math/prim/constraint/cholesky_corr_constrain.hpp

@@ -43,9 +43,11 @@ cholesky_corr_constrain(const EigVec& y, int K) {
 }
 
 // FIXME to match above after debugged
-template <typename EigVec, require_eigen_vector_t<EigVec>* = nullptr>
+template <typename EigVec, typename Lp,
+          require_eigen_vector_t<EigVec>* = nullptr,
+          require_convertible_t<return_type_t<EigVec>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<EigVec>, Eigen::Dynamic, Eigen::Dynamic>
-cholesky_corr_constrain(const EigVec& y, int K, return_type_t<EigVec>& lp) {
+cholesky_corr_constrain(const EigVec& y, int K, Lp& lp) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
   using std::sqrt;
@@ -85,13 +87,17 @@ cholesky_corr_constrain(const EigVec& y, int K, return_type_t<EigVec>& lp) {
  * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
  *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
  *  and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param y Linearly Serialized vector of size `(K * (K - 1))/2` holding the
  *  column major order elements of the lower triangurlar
  * @param K The size of the matrix to return
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_not_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_corr_constrain(const T& y, int K, Lp& lp) {
   if (Jacobian) {
     return cholesky_corr_constrain(y, K, lp);
   } else {
@@ -110,13 +116,17 @@ inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
  * @tparam T A standard vector with inner type inheriting from
  * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
  * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param y Linearly Serialized vector of size `(K * (K - 1))/2` holding the
  *  column major order elements of the lower triangurlar
  * @param K The size of the matrix to return
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto cholesky_corr_constrain(const T& y, int K, return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_corr_constrain(const T& y, int K, Lp& lp) {
   return apply_vector_unary<T>::apply(y, [&lp, K](auto&& v) {
     return cholesky_corr_constrain<Jacobian>(v, K, lp);
   });

stan/math/prim/constraint/cholesky_factor_constrain.hpp

@@ -70,9 +70,10 @@ cholesky_factor_constrain(const T& x, int M, int N) {
  * determinant
  * @return Cholesky factor
  */
-template <typename T, require_eigen_vector_t<T>* = nullptr>
+template <typename T, typename Lp, require_eigen_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
-cholesky_factor_constrain(const T& x, int M, int N, return_type_t<T>& lp) {
+cholesky_factor_constrain(const T& x, int M, int N, Lp& lp) {
   check_size_match("cholesky_factor_constrain", "x.size()", x.size(),
                    "((N * (N + 1)) / 2 + (M - N) * N)",
                    ((N * (N + 1)) / 2 + (M - N) * N));
@@ -99,15 +100,18 @@ cholesky_factor_constrain(const T& x, int M, int N, return_type_t<T>& lp) {
  * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
  *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
  *  and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained values
  * @param M number of rows
  * @param N number of columns
  * @param[in,out] lp log density accumulator
  * @return Cholesky factor
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto cholesky_factor_constrain(const T& x, int M, int N,
-                                      return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_not_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_factor_constrain(const T& x, int M, int N, Lp& lp) {
   if (Jacobian) {
     return cholesky_factor_constrain(x, M, N, lp);
   } else {
@@ -129,15 +133,18 @@ inline auto cholesky_factor_constrain(const T& x, int M, int N,
  * @tparam T A standard vector with inner type inheriting from
  * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
  * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained values
  * @param M number of rows
  * @param N number of columns
  * @param[in,out] lp log density accumulator
  * @return Cholesky factor
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto cholesky_factor_constrain(const T& x, int M, int N,
-                                      return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cholesky_factor_constrain(const T& x, int M, int N, Lp& lp) {
   return apply_vector_unary<T>::apply(x, [&lp, M, N](auto&& v) {
     return cholesky_factor_constrain<Jacobian>(v, M, N, lp);
   });

stan/math/prim/constraint/corr_matrix_constrain.hpp

@@ -61,13 +61,16 @@ corr_matrix_constrain(const T& x, Eigen::Index k) {
  *
  * @tparam T type of the vector (must be derived from \c Eigen::MatrixBase and
  * have one compile-time dimension equal to 1)
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained partial correlations.
  * @param k Dimensionality of returned correlation matrix.
  * @param lp Log probability reference to increment.
  */
-template <typename T, require_eigen_col_vector_t<T>* = nullptr>
+template <typename T, typename Lp, require_eigen_col_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
-corr_matrix_constrain(const T& x, Eigen::Index k, return_type_t<T>& lp) {
+corr_matrix_constrain(const T& x, Eigen::Index k, Lp& lp) {
   Eigen::Index k_choose_2 = (k * (k - 1)) / 2;
   check_size_match("cov_matrix_constrain", "x.size()", x.size(), "k_choose_2",
                    k_choose_2);
@@ -89,13 +92,16 @@ corr_matrix_constrain(const T& x, Eigen::Index k, return_type_t<T>& lp) {
  * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
  *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
  *  and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x Vector of unconstrained partial correlations
  * @param k Dimensionality of returned correlation matrix
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto corr_matrix_constrain(const T& x, Eigen::Index k,
-                                  return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_not_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto corr_matrix_constrain(const T& x, Eigen::Index k, Lp& lp) {
   if (Jacobian) {
     return corr_matrix_constrain(x, k, lp);
   } else {
@@ -118,12 +124,16 @@ inline auto corr_matrix_constrain(const T& x, Eigen::Index k,
  * @tparam T A standard vector with inner type inheriting from
  * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
  * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param y Vector of unconstrained partial correlations
  * @param K Dimensionality of returned correlation matrix
  * @param[in,out] lp log density accumulator
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto corr_matrix_constrain(const T& y, int K, return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto corr_matrix_constrain(const T& y, int K, Lp& lp) {
   return apply_vector_unary<T>::apply(y, [&lp, K](auto&& v) {
     return corr_matrix_constrain<Jacobian>(v, K, lp);
   });

stan/math/prim/constraint/cov_matrix_constrain.hpp

@@ -55,14 +55,17 @@ cov_matrix_constrain(const T& x, Eigen::Index K) {
  *
  * @tparam T type of the vector (must be derived from \c Eigen::MatrixBase and
  * have one compile-time dimension equal to 1)
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x The vector to convert to a covariance matrix.
  * @param K The dimensions of the resulting covariance matrix.
  * @param lp Reference
  * @throws std::domain_error if (x.size() != K + (K choose 2)).
  */
-template <typename T, require_eigen_col_vector_t<T>* = nullptr>
+template <typename T, typename Lp, require_eigen_col_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
 inline Eigen::Matrix<value_type_t<T>, Eigen::Dynamic, Eigen::Dynamic>
-cov_matrix_constrain(const T& x, Eigen::Index K, return_type_t<T>& lp) {
+cov_matrix_constrain(const T& x, Eigen::Index K, Lp& lp) {
   using Eigen::Dynamic;
   using Eigen::Matrix;
   using std::exp;
@@ -99,14 +102,17 @@ cov_matrix_constrain(const T& x, Eigen::Index K, return_type_t<T>& lp) {
  * @tparam T A type inheriting from `Eigen::DenseBase` or a `var_value` with
  *  inner type inheriting from `Eigen::DenseBase` with compile time dynamic rows
  *  and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x The vector to convert to a covariance matrix
  * @param K The dimensions of the resulting covariance matrix
  * @param[in, out] lp log density accumulator
  * @throws std::domain_error if (x.size() != K + (K choose 2)).
  */
-template <bool Jacobian, typename T, require_not_std_vector_t<T>* = nullptr>
-inline auto cov_matrix_constrain(const T& x, Eigen::Index K,
-                                 return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_not_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cov_matrix_constrain(const T& x, Eigen::Index K, Lp& lp) {
   if (Jacobian) {
     return cov_matrix_constrain(x, K, lp);
   } else {
@@ -127,14 +133,17 @@ inline auto cov_matrix_constrain(const T& x, Eigen::Index K,
  * @tparam T A standard vector with inner type inheriting from
  * `Eigen::DenseBase` or a `var_value` with inner type inheriting from
  * `Eigen::DenseBase` with compile time dynamic rows and 1 column
+ * @tparam Lp A scalar type for the lp argument. The scalar type of T should be
+ * convertable to this.
  * @param x The vector to convert to a covariance matrix
  * @param K The dimensions of the resulting covariance matrix
  * @param[in, out] lp log density accumulator
  * @throws std::domain_error if (x.size() != K + (K choose 2)).
  */
-template <bool Jacobian, typename T, require_std_vector_t<T>* = nullptr>
-inline auto cov_matrix_constrain(const T& x, Eigen::Index K,
-                                 return_type_t<T>& lp) {
+template <bool Jacobian, typename T, typename Lp,
+          require_std_vector_t<T>* = nullptr,
+          require_convertible_t<return_type_t<T>, Lp>* = nullptr>
+inline auto cov_matrix_constrain(const T& x, Eigen::Index K, Lp& lp) {
   return apply_vector_unary<T>::apply(x, [&lp, K](auto&& v) {
     return cov_matrix_constrain<Jacobian>(v, K, lp);
   });

stan/math/prim/constraint/lb_constrain.hpp

@@ -50,14 +50,17 @@ inline auto lb_constrain(const T& x, const L& lb) {
  *
  * @tparam T Scalar.
  * @tparam L Scalar.
+ * @tparam Lp scalar, T and L should be convertible to this.
  * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
-template <typename T, typename L, require_all_stan_scalar_t<T, L>* = nullptr,
-          require_all_not_st_var<T, L>* = nullptr>
-inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
+template <typename T, typename L, typename Lp,
+          require_all_stan_scalar_t<T, L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr,
+          require_convertible_t<return_type_t<T, L>, Lp>* = nullptr>
+inline auto lb_constrain(const T& x, const L& lb, Lp& lp) {
   if (value_of_rec(lb) == NEGATIVE_INFTY) {
     return identity_constrain(x, lb);
   } else {
@@ -89,15 +92,18 @@ inline auto lb_constrain(T&& x, L&& lb) {
  *
  * @tparam T A type inheriting from `EigenBase`.
  * @tparam L Scalar.
+ * @tparam Lp scalar, T and L should be convertible to this.
  * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
-template <typename T, typename L, require_eigen_t<T>* = nullptr,
+template <typename T, typename L, typename Lp, require_eigen_t<T>* = nullptr,
           require_stan_scalar_t<L>* = nullptr,
-          require_all_not_st_var<T, L>* = nullptr>
-inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
+          require_all_not_st_var<T, L>* = nullptr,
+          require_convertible_t<return_type_t<T, L>, Lp>* = nullptr>
+
+inline auto lb_constrain(const T& x, const L& lb, Lp& lp) {
   return eval(
       x.unaryExpr([lb, &lp](auto&& xx) { return lb_constrain(xx, lb, lp); }));
 }
@@ -126,14 +132,17 @@ inline auto lb_constrain(T&& x, L&& lb) {
  *
  * @tparam T A type inheriting from `EigenBase`.
  * @tparam L A type inheriting from `EigenBase`.
+ * @tparam Lp Scalar, the scalar type of T and L should be convertible to this.
  * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
-template <typename T, typename L, require_all_eigen_t<T, L>* = nullptr,
-          require_all_not_st_var<T, L>* = nullptr>
-inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
+template <typename T, typename L, typename Lp,
+          require_all_eigen_t<T, L>* = nullptr,
+          require_all_not_st_var<T, L>* = nullptr,
+          require_convertible_t<return_type_t<T, L>, Lp>* = nullptr>
+inline auto lb_constrain(const T& x, const L& lb, Lp& lp) {
   check_matching_dims("lb_constrain", "x", x, "lb", lb);
   return eval(x.binaryExpr(
       lb, [&lp](auto&& xx, auto&& lbb) { return lb_constrain(xx, lbb, lp); }));
@@ -164,14 +173,16 @@ inline auto lb_constrain(const std::vector<T>& x, const L& lb) {
  *
  * @tparam T A Any type with a Scalar `scalar_type`.
  * @tparam L A type inheriting from `EigenBase` or a standard vector.
+ * @tparam Lp Scalar, T and the scalar type of L should be convertible to this.
  * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
-template <typename T, typename L, require_not_std_vector_t<L>* = nullptr>
-inline auto lb_constrain(const std::vector<T>& x, const L& lb,
-                         return_type_t<T, L>& lp) {
+template <typename T, typename L, typename Lp,
+          require_not_std_vector_t<L>* = nullptr,
+          require_convertible_t<return_type_t<T, L>, Lp>* = nullptr>
+inline auto lb_constrain(const std::vector<T>& x, const L& lb, Lp& lp) {
   std::vector<plain_type_t<decltype(lb_constrain(x[0], lb))>> ret(x.size());
   for (size_t i = 0; i < x.size(); ++i) {
     ret[i] = lb_constrain(x[i], lb, lp);
@@ -205,14 +216,16 @@ inline auto lb_constrain(const std::vector<T>& x, const std::vector<L>& lb) {
  *
  * @tparam T A Any type with a Scalar `scalar_type`.
  * @tparam L A type inheriting from `EigenBase` or a standard vector.
+ * @tparam Lp Scalar, T and the scalar type of L should be convertible to this.
  * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in,out] lp reference to log probability to increment
  * @return lower-bound constrained value corresponding to inputs
  */
-template <typename T, typename L>
+template <typename T, typename L, typename Lp,
+          require_convertible_t<return_type_t<T, L>, Lp>* = nullptr>
 inline auto lb_constrain(const std::vector<T>& x, const std::vector<L>& lb,
-                         return_type_t<T, L>& lp) {
+                         Lp& lp) {
   check_matching_dims("lb_constrain", "x", x, "lb", lb);
   std::vector<plain_type_t<decltype(lb_constrain(x[0], lb[0]))>> ret(x.size());
   for (size_t i = 0; i < x.size(); ++i) {
@@ -234,13 +247,15 @@ inline auto lb_constrain(const std::vector<T>& x, const std::vector<L>& lb,
  * type inheriting from `Eigen::EigenBase`, a standard vector, or a scalar
  * @tparam L A type inheriting from `Eigen::EigenBase`, a `var_value` with inner
  * type inheriting from `Eigen::EigenBase`, a standard vector, or a scalar
+ * @tparam Lp Scalar, the scalar type of T and L should be convertible to this.
  * @param[in] x unconstrained input
  * @param[in] lb lower bound on output
  * @param[in, out] lp log density accumulator
  * @return lower-bound constrained value corresponding to inputs
  */
-template <bool Jacobian, typename T, typename L>
-inline auto lb_constrain(const T& x, const L& lb, return_type_t<T, L>& lp) {
+template <bool Jacobian, typename T, typename L, typename Lp,
+          require_convertible_t<return_type_t<T, L>, Lp>* = nullptr>
+inline auto lb_constrain(const T& x, const L& lb, Lp& lp) {
   if (Jacobian) {
     return lb_constrain(x, lb, lp);
   } else {