Merge pull request #3221 from stan-dev/fix/exprs-apply

Use Perfect Forwarding in all functions that use `apply` family of functors

Author: SteveBronder

stan/math/fwd/fun/inv_logit.hpp

@@ -16,10 +16,10 @@ namespace math {
  * @param x argument
  * @return inverse logit of argument
  */
-template <typename T>
-inline fvar<T> inv_logit(const fvar<T>& x) {
-  return fvar<T>(inv_logit(x.val_),
-                 x.d_ * inv_logit(x.val_) * (1 - inv_logit(x.val_)));
+template <typename T, require_fvar_t<T>* = nullptr>
+inline auto inv_logit(T&& x) {
+  return std::decay_t<T>(inv_logit(x.val_),
+                         x.d_ * inv_logit(x.val_) * (1 - inv_logit(x.val_)));
 }
 
 }  // namespace math

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 {
@@ -19,24 +20,24 @@ namespace math {
  * @throw std::domain_error If the input vector is size 0.
  */
 template <typename T, require_vector_st<is_fvar, T>* = nullptr>
-inline auto log_softmax(const T& x) {
-  return apply_vector_unary<T>::apply(x, [&](const auto& alpha) {
+inline auto log_softmax(T&& x) {
+  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

@@ -50,9 +50,9 @@ inline fvar<T> log_sum_exp(const fvar<T>& x1, double x2) {
  * @return The log of the sum of the exponentiated vector values.
  */
 template <typename T, require_container_st<is_fvar, T>* = nullptr>
-inline auto log_sum_exp(const T& x) {
+inline auto log_sum_exp(T&& x) {
   return apply_vector_unary<ref_type_t<T>>::reduce(
-      to_ref(x), [&](const 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

@@ -15,14 +15,14 @@ namespace math {
 /**
  * Compute the L1 norm of the specified vector of values.
  *
- * @tparam T Type of input vector.
+ * @tparam Container Type of input vector.
  * @param[in] x Vector of specified values.
  * @return L1 norm of x.
  */
 template <typename Container, require_eigen_vt<is_fvar, Container>* = nullptr>
-inline auto norm1(const Container& x) {
+inline auto norm1(Container&& x) {
   return apply_vector_unary<ref_type_t<Container>>::reduce(
-      to_ref(x), [&](const 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

@@ -14,14 +14,14 @@ namespace math {
 /**
  * Compute the L2 norm of the specified vector of values.
  *
- * @tparam T Type of input vector.
+ * @tparam Container Type of input vector.
  * @param[in] x Vector of specified values.
  * @return L2 norm of x.
  */
 template <typename Container, require_eigen_vt<is_fvar, Container>* = nullptr>
-inline auto norm2(const Container& x) {
+inline auto norm2(Container&& x) {
   return apply_vector_unary<ref_type_t<Container>>::reduce(
-      to_ref(x), [&](const 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/fwd/fun/pow.hpp

@@ -73,9 +73,13 @@ inline auto pow(const T1& x1, const T2& x2) {
 template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr,
           require_all_not_matrix_st<is_var, T1, T2>* = nullptr,
           require_any_fvar_t<base_type_t<T1>, base_type_t<T2>>* = nullptr>
-inline auto pow(const T1& a, const T2& b) {
+inline auto pow(T1&& a, T2&& b) {
   return apply_scalar_binary(
-      [](const auto& c, const auto& d) { return stan::math::pow(c, d); }, a, b);
+      [](auto&& c, auto&& d) {
+        return stan::math::pow(std::forward<decltype(c)>(c),
+                               std::forward<decltype(d)>(d));
+      },
+      std::forward<T1>(a), std::forward<T2>(b));
 }
 
 }  // namespace math

stan/math/fwd/functor/apply_scalar_unary.hpp

@@ -18,20 +18,23 @@ namespace math {
  * autodiff variable.
  */
 template <typename F, typename T>
-struct apply_scalar_unary<F, fvar<T> > {
+struct apply_scalar_unary<F, T, require_fvar_t<T>> {
   /**
    * Function return type, which is same as the argument type for
    * the function, <code>fvar&lt;T&gt;</code>.
    */
-  using return_t = fvar<T>;
+  using return_t = std::decay_t<T>;
 
   /**
    * Apply the function specified by F to the specified argument.
    *
    * @param x Argument variable.
    * @return Function applied to the variable.
    */
-  static inline return_t apply(const fvar<T>& x) { return F::fun(x); }
+  template <typename T2>
+  static inline auto apply(T2&& x) {
+    return F::fun(std::forward<T2>(x));
+  }
 };
 
 }  // namespace math

stan/math/fwd/functor/finite_diff.hpp

@@ -44,8 +44,8 @@ inline constexpr double aggregate_tangent(const FuncTangent& tangent,
 template <typename FuncTangent, typename InputArg,
           require_st_fvar<InputArg>* = nullptr>
 inline auto aggregate_tangent(const FuncTangent& tangent, const InputArg& arg) {
-  return sum(apply_scalar_binary(
-      [](const auto& x, const auto& y) { return x * y.d_; }, tangent, arg));
+  return sum(apply_scalar_binary([](auto&& x, auto&& y) { return x * y.d_; },
+                                 tangent, arg));
 }
 }  // namespace internal
 
@@ -73,7 +73,7 @@ inline auto finite_diff(const F& func, const TArgs&... args) {
   std::vector<FvarInnerT> serialised_args
       = serialize<FvarInnerT>(value_of(args)...);
 
-  auto serial_functor = [&](const auto& v) {
+  auto serial_functor = [&](auto&& v) {
     auto v_deserializer = to_deserializer(v);
     return func(v_deserializer.read(args)...);
   };

stan/math/opencl/prim/dirichlet_lpdf.hpp

@@ -2,6 +2,7 @@
 #define STAN_MATH_OPENCL_PRIM_DIRICHLET_LPDF_HPP
 #ifdef STAN_OPENCL
 
+#include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/constants.hpp>

stan/math/opencl/prim/log_softmax.hpp

@@ -22,8 +22,7 @@ template <typename T,
           require_all_kernel_expressions_and_none_scalar_t<T>* = nullptr>
 inline matrix_cl<double> log_softmax(const T& a) {
   check_nonzero_size("log_softmax (OpenCL)", "x", a);
-  return make_holder_cl([](const auto& x) { return x - log_sum_exp(x); },
-                        to_ref(a));
+  return make_holder_cl([](auto&& x) { return x - log_sum_exp(x); }, to_ref(a));
 }
 
 }  // namespace math