use a seperate class for csr_matrix adjoint

Author: SteveBronder

stan/math/rev/fun/csr_matrix_times_vector.hpp

@@ -11,6 +11,50 @@
 namespace stan {
 namespace math {
 
+namespace internal {
+  template <typename Result_, typename WMat_, typename B_>
+  struct csr_adjoint : public vari {
+    std::decay_t<Result_> res_;
+    std::decay_t<WMat_> w_mat_;
+    std::decay_t<B_> b_;
+    template <typename T1, typename T2, typename T3>
+    csr_adjoint(T1&& res, T2&& w_mat, T3&& b)
+        : vari(0.0), res_(std::forward<T1>(res)),
+          w_mat_(std::forward<T2>(w_mat)), b_(std::forward<T3>(b)) {}
+
+    void chain() {
+      chain_internal(res_, w_mat_, b_);
+    }
+    template <typename Result, typename WMat, typename B,
+      require_rev_matrix_t<WMat>* = nullptr,
+      require_rev_matrix_t<B>* = nullptr>
+    void chain_internal(Result&& res, WMat&& w_mat, B&& b) {
+      w_mat.adj() += res.adj() * b.val().transpose();
+      b.adj() += w_mat.val().transpose() * res.adj();
+    }
+
+    template <typename Result, typename WMat, typename B,
+      require_rev_matrix_t<WMat>* = nullptr,
+      require_not_rev_matrix_t<B>* = nullptr>
+    void chain_internal(Result&& res, WMat&& w_mat, B&& b) {
+      w_mat.adj() += res.adj() * b.transpose();
+    }
+
+    template <typename Result, typename WMat, typename B,
+      require_not_rev_matrix_t<WMat>* = nullptr,
+      require_rev_matrix_t<B>* = nullptr>
+    void chain_internal(Result&& res, WMat&& w_mat, B&& b) {
+      b.adj() += w_mat.transpose() * res.adj();
+    }
+  };
+  template <typename Result_, typename WMat_, typename B_>
+  inline vari* make_csr_adjoint(Result_&& res, WMat_&& w_mat, B_&& b) {
+    return new csr_adjoint<Result_, WMat_, B_>(
+        std::forward<Result_>(res), std::forward<WMat_>(w_mat),
+        std::forward<B_>(b));
+  }
+}
+
 /**
  * \addtogroup csr_format
  * Return the multiplication of the sparse matrix (specified by
@@ -74,29 +118,22 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
     sparse_var_value_t w_mat_arena
         = to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w, u_arena, v_arena);
     arena_t<return_t> res = w_mat_arena.val() * value_of(b_arena);
-    reverse_pass_callback([res, w_mat_arena, b_arena]() mutable {
-      w_mat_arena.adj() += res.adj() * b_arena.val().transpose();
-      b_arena.adj() += w_mat_arena.val().transpose() * res.adj();
-    });
+    stan::math::internal::make_csr_adjoint(res, w_mat_arena, b_arena);
     return return_t(res);
   } else if (!is_constant<T2>::value) {
     arena_t<promote_scalar_t<var, T2>> b_arena = b;
     auto w_val_arena = to_arena(value_of(w));
     sparse_val_mat w_val_mat(m, n, w_val_arena.size(), u_arena.data(),
                              v_arena.data(), w_val_arena.data());
     arena_t<return_t> res = w_val_mat * value_of(b_arena);
-    reverse_pass_callback([w_val_mat, res, b_arena]() mutable {
-      b_arena.adj() += w_val_mat.transpose() * res.adj();
-    });
+    stan::math::internal::make_csr_adjoint(res, w_val_mat, b_arena);
     return return_t(res);
   } else {
     sparse_var_value_t w_mat_arena
         = to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w, u_arena, v_arena);
     auto b_arena = to_arena(value_of(b));
     arena_t<return_t> res = w_mat_arena.val() * b_arena;
-    reverse_pass_callback([res, w_mat_arena, b_arena]() mutable {
-      w_mat_arena.adj() += res.adj() * b_arena.transpose();
-    });
+    stan::math::internal::make_csr_adjoint(res, w_mat_arena, b_arena);
     return return_t(res);
   }
 }