update with brians review

Author: SteveBronder

stan/math/fwd/fun/log_softmax.hpp

@@ -6,6 +6,7 @@
 #include <stan/math/fwd/meta.hpp>
 #include <stan/math/fwd/fun/softmax.hpp>
 #include <stan/math/prim/fun/log_softmax.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
 
 namespace stan {
 namespace math {
@@ -20,23 +21,23 @@ namespace math {
  */
 template <typename T, require_vector_st<is_fvar, T>* = nullptr>
 inline auto log_softmax(T&& x) {
-  return apply_vector_unary<T>::apply(std::forward<T>(x), [&](auto&& alpha) {
+  return apply_vector_unary<T>::apply(std::forward<T>(x), [](auto&& alpha) {
     using T_alpha = decltype(alpha);
     using T_fvar = value_type_t<T_alpha>;
     using T_fvar_inner = typename T_fvar::Scalar;
 
-    const Eigen::Ref<const plain_type_t<T_alpha>>& alpha_ref = alpha;
+    auto&& alpha_ref = to_ref(std::forward<decltype(alpha)>(alpha));
     Eigen::Matrix<T_fvar_inner, -1, 1> alpha_t = alpha_ref.val();
     Eigen::Matrix<T_fvar_inner, -1, 1> softmax_alpha_t = softmax(alpha_t);
 
-    Eigen::Matrix<T_fvar, -1, 1> log_softmax_alpha(alpha.size());
+    Eigen::Matrix<T_fvar, -1, 1> log_softmax_alpha(alpha_ref.size());
     log_softmax_alpha.val() = log_softmax(alpha_t);
     log_softmax_alpha.d().setZero();
 
-    for (int m = 0; m < alpha.size(); ++m) {
+    for (int m = 0; m < alpha_ref.size(); ++m) {
       T_fvar_inner negative_alpha_m_d_times_softmax_alpha_t_m
           = -alpha_ref.coeff(m).d_ * softmax_alpha_t(m);
-      for (int k = 0; k < alpha.size(); ++k) {
+      for (int k = 0; k < alpha_ref.size(); ++k) {
         if (m == k) {
           log_softmax_alpha(k).d_
               += alpha_ref.coeff(m).d_

stan/math/fwd/fun/log_sum_exp.hpp

@@ -52,7 +52,7 @@ inline fvar<T> log_sum_exp(const fvar<T>& x1, double x2) {
 template <typename T, require_container_st<is_fvar, T>* = nullptr>
 inline auto log_sum_exp(T&& x) {
   return apply_vector_unary<ref_type_t<T>>::reduce(
-      to_ref(std::forward<T>(x)), [&](auto&& v) {
+      to_ref(std::forward<T>(x)), [](auto&& v) {
         using T_fvar_inner = typename value_type_t<decltype(v)>::Scalar;
         using mat_type = Eigen::Matrix<T_fvar_inner, -1, -1>;
         mat_type vals = v.val();

stan/math/fwd/fun/norm1.hpp

@@ -22,7 +22,7 @@ namespace math {
 template <typename Container, require_eigen_vt<is_fvar, Container>* = nullptr>
 inline auto norm1(Container&& x) {
   return apply_vector_unary<ref_type_t<Container>>::reduce(
-      to_ref(std::forward<Container>(x)), [&](auto&& v) {
+      to_ref(std::forward<Container>(x)), [](auto&& v) {
         using T_fvar_inner = typename value_type_t<decltype(v)>::Scalar;
         return fvar<T_fvar_inner>(norm1(v.val()),
                                   v.d().cwiseProduct(sign(v.val())).sum());

stan/math/fwd/fun/norm2.hpp

@@ -21,7 +21,7 @@ namespace math {
 template <typename Container, require_eigen_vt<is_fvar, Container>* = nullptr>
 inline auto norm2(Container&& x) {
   return apply_vector_unary<ref_type_t<Container>>::reduce(
-      to_ref(std::forward<Container>(x)), [&](auto&& v) {
+      to_ref(std::forward<Container>(x)), [](auto&& v) {
         using T_fvar_inner = typename value_type_t<decltype(v)>::Scalar;
         T_fvar_inner res = norm2(v.val());
         return fvar<T_fvar_inner>(res,

stan/math/prim/constraint/prob_constrain.hpp

@@ -49,7 +49,7 @@ inline auto prob_constrain(T&& x) {
  */
 template <typename T, typename Lp>
 inline auto prob_constrain(T&& x, Lp& lp) {
-  std::decay_t<T> log_inv_logit_x = log_inv_logit(x);
+  plain_type_t<T> log_inv_logit_x = log_inv_logit(x);
   lp += log_inv_logit_x + log1m_inv_logit(std::forward<T>(x));
   return exp(std::move(log_inv_logit_x));
 }

stan/math/prim/fun/abs.hpp

@@ -80,7 +80,7 @@ template <typename Container,
           require_container_bt<std::is_arithmetic, Container>* = nullptr>
 inline auto abs(Container&& x) {
   return apply_vector_unary<Container>::apply(
-      std::forward<Container>(x), [&](auto&& v) { return v.array().abs(); });
+      std::forward<Container>(x), [](auto&& v) { return v.array().abs(); });
 }
 
 namespace internal {

stan/math/prim/fun/as_array_or_scalar.hpp

@@ -59,7 +59,13 @@ inline auto as_array_or_scalar(T&& v) {
 
 /**
  * Converts a std::vector type to an array.
- *
+ * @note The math library's reverse mode assumes that `Eigen::Map`
+ *  types are allocated and owned elsewhere so we cannot just return
+ *  back a map here else the reverse mode library
+ *  may try to access into a dangling pointer. Instead we wrap
+ *  the `Eigen::Map` in a `Holder` to trick the reverse mode library
+ *  into not thinking this is a map. The `.array().matrix()` inside the
+ *  holder is so that the holder thinks it is returning an expression.
  * @tparam T Type of scalar element.
  * @param v Specified vector.
  * @return Matrix converted to an array.

stan/math/prim/fun/cbrt.hpp

@@ -40,7 +40,7 @@ template <
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr,
     require_container_t<T>* = nullptr>
 inline auto cbrt(T&& x) {
-  return apply_scalar_unary<cbrt_fun, T>::apply(x);
+  return apply_scalar_unary<cbrt_fun, T>::apply(std::forward<T>(x));
 }
 
 }  // namespace math

stan/math/prim/fun/cos.hpp

@@ -77,7 +77,7 @@ template <typename Container,
           require_container_bt<std::is_arithmetic, Container>* = nullptr>
 inline auto cos(Container&& x) {
   return apply_vector_unary<Container>::apply(
-      std::forward<Container>(x), [&](auto&& v) { return v.array().cos(); });
+      std::forward<Container>(x), [](auto&& v) { return v.array().cos(); });
 }
 
 namespace internal {

stan/math/prim/fun/inv_square.hpp

@@ -27,7 +27,9 @@ template <typename Container,
               Container>* = nullptr,
           require_container_t<Container>* = nullptr>
 inline auto inv_square(Container&& x) {
-  return inv(square(std::forward<Container>(x)));
+  return make_holder([](auto&& v) { 
+    return inv(square(std::forward<decltype(v)>(v))); 
+  }, std::forward<Container>(x));
 }
 
 /**