From 8b96c450cce144caf71ffca11cd54915d3205711 Mon Sep 17 00:00:00 2001
From: stevebronder <stevebronder@stevebronder.com>
Date: Fri, 12 Apr 2024 17:22:07 -0400
Subject: [PATCH 1/9] fixes aos csr matrix bug, still debugging soa matrix bug

---
 stan/math/prim/fun/value_of.hpp               | 24 +++++-
 stan/math/prim/meta/promote_scalar_type.hpp   | 11 ++-
 stan/math/rev/fun.hpp                         |  1 +
 stan/math/rev/fun/csr_matrix_times_vector.hpp | 68 ++++++---------
 stan/math/rev/fun/to_soa_sparse_matrix.hpp    | 84 +++++++++++++++++++
 stan/math/rev/fun/to_var_value.hpp            |  6 +-
 .../mix/fun/csr_matrix_times_vector_test.cpp  | 21 ++++-
 7 files changed, 163 insertions(+), 52 deletions(-)
 create mode 100644 stan/math/rev/fun/to_soa_sparse_matrix.hpp

diff --git a/stan/math/prim/fun/value_of.hpp b/stan/math/prim/fun/value_of.hpp
index 7cc37ac9b2c..cdb5115d104 100644
--- a/stan/math/prim/fun/value_of.hpp
+++ b/stan/math/prim/fun/value_of.hpp
@@ -1,6 +1,7 @@
 #ifndef STAN_MATH_PRIM_FUN_VALUE_OF_HPP
 #define STAN_MATH_PRIM_FUN_VALUE_OF_HPP
 
+#include <iostream>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <cstddef>
@@ -67,7 +68,7 @@ inline auto value_of(const T& x) {
  * @param[in] M Matrix to be converted
  * @return Matrix of values
  **/
-template <typename EigMat, require_eigen_t<EigMat>* = nullptr,
+template <typename EigMat, require_eigen_dense_base_t<EigMat>* = nullptr,
           require_not_st_arithmetic<EigMat>* = nullptr>
 inline auto value_of(EigMat&& M) {
   return make_holder(
@@ -77,6 +78,27 @@ inline auto value_of(EigMat&& M) {
       std::forward<EigMat>(M));
 }
 
+template <typename EigMat, require_eigen_sparse_base_t<EigMat>* = nullptr,
+          require_not_st_arithmetic<EigMat>* = nullptr>
+inline auto value_of(EigMat&& M) {
+  auto&& M_ref = to_ref(M);
+  using scalar_t = decltype(value_of(std::declval<value_type_t<EigMat>>()));
+  promote_scalar_t<scalar_t, plain_type_t<EigMat>> ret(M_ref.rows(), M_ref.cols());
+  ret.reserve(M_ref.nonZeros());
+  for (int k = 0; k < M_ref.outerSize(); ++k) {
+    for (typename std::decay_t<EigMat>::InnerIterator it(M_ref, k); it; ++it) {
+      ret.insert(it.row(), it.col()) = value_of(it.valueRef());
+    }
+  }
+  ret.makeCompressed();
+  return ret;
+}
+template <typename EigMat, require_eigen_sparse_base_t<EigMat>* = nullptr,
+          require_st_arithmetic<EigMat>* = nullptr>
+inline auto value_of(EigMat&& M) {
+  return std::forward<EigMat>(M);
+}
+
 }  // namespace math
 }  // namespace stan
 
diff --git a/stan/math/prim/meta/promote_scalar_type.hpp b/stan/math/prim/meta/promote_scalar_type.hpp
index fee623ab141..0c30b5a6140 100644
--- a/stan/math/prim/meta/promote_scalar_type.hpp
+++ b/stan/math/prim/meta/promote_scalar_type.hpp
@@ -80,7 +80,7 @@ struct promote_scalar_type<T, S,
  * @tparam S input matrix type
  */
 template <typename T, typename S>
-struct promote_scalar_type<T, S, require_eigen_t<S>> {
+struct promote_scalar_type<T, S, require_eigen_dense_base_t<S>> {
   /**
    * The promoted type.
    */
@@ -93,6 +93,15 @@ struct promote_scalar_type<T, S, require_eigen_t<S>> {
                    S::RowsAtCompileTime, S::ColsAtCompileTime>>::type;
 };
 
+template <typename T, typename S>
+struct promote_scalar_type<T, S, require_eigen_sparse_base_t<S>> {
+  /**
+   * The promoted type.
+   */
+  using type = Eigen::SparseMatrix<typename promote_scalar_type<T, typename S::Scalar>::type,
+   S::Options, typename S::StorageIndex>;
+};
+
 template <typename... PromotionScalars, typename... UnPromotedTypes>
 struct promote_scalar_type<std::tuple<PromotionScalars...>,
                            std::tuple<UnPromotedTypes...>> {
diff --git a/stan/math/rev/fun.hpp b/stan/math/rev/fun.hpp
index 9324df2edcc..2bba861f771 100644
--- a/stan/math/rev/fun.hpp
+++ b/stan/math/rev/fun.hpp
@@ -183,6 +183,7 @@
 #include <stan/math/rev/fun/tgamma.hpp>
 #include <stan/math/rev/fun/to_var.hpp>
 #include <stan/math/rev/fun/to_arena.hpp>
+#include <stan/math/rev/fun/to_soa_sparse_matrix.hpp>
 #include <stan/math/rev/fun/to_var_value.hpp>
 #include <stan/math/rev/fun/to_vector.hpp>
 #include <stan/math/rev/fun/trace.hpp>
diff --git a/stan/math/rev/fun/csr_matrix_times_vector.hpp b/stan/math/rev/fun/csr_matrix_times_vector.hpp
index d0cb4bd2373..8857f1ce23d 100644
--- a/stan/math/rev/fun/csr_matrix_times_vector.hpp
+++ b/stan/math/rev/fun/csr_matrix_times_vector.hpp
@@ -2,6 +2,7 @@
 #define STAN_MATH_REV_FUN_CSR_MATRIX_TIMES_VECTOR_HPP
 
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/rev/fun/to_soa_sparse_matrix.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/csr_u_to_z.hpp>
@@ -11,35 +12,21 @@ namespace stan {
 namespace math {
 
 namespace internal {
-template <typename T1, typename T2, typename Res,
-          require_eigen_t<T1>* = nullptr>
-void update_w(T1& w, int m, int n, std::vector<int, arena_allocator<int>>& u,
-              std::vector<int, arena_allocator<int>>& v, T2&& b, Res&& res) {
-  Eigen::Map<Eigen::SparseMatrix<var, Eigen::RowMajor>> w_mat(
-      m, n, w.size(), u.data(), v.data(), w.data());
-  for (int k = 0; k < w_mat.outerSize(); ++k) {
-    for (Eigen::Map<Eigen::SparseMatrix<var, Eigen::RowMajor>>::InnerIterator
-             it(w_mat, k);
-         it; ++it) {
-      it.valueRef().adj()
-          += res.adj().coeff(it.row()) * value_of(b).coeff(it.col());
-    }
-  }
-}
 
 template <typename T1, typename T2, typename Res,
           require_var_matrix_t<T1>* = nullptr>
-void update_w(T1& w, int m, int n, std::vector<int, arena_allocator<int>>& u,
-              std::vector<int, arena_allocator<int>>& v, T2&& b, Res&& res) {
-  Eigen::Map<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat(
-      m, n, w.size(), u.data(), v.data(), w.adj().data());
-  for (int k = 0; k < w_mat.outerSize(); ++k) {
-    for (Eigen::Map<Eigen::SparseMatrix<double, Eigen::RowMajor>>::InnerIterator
-             it(w_mat, k);
+inline void update_w(T1& w, T2&& b, Res&& res) {
+      std::cout << "pre w adj: \n" << w.adj() << std::endl;
+      std::cout << "pre b: \n" << value_of(b).eval() << std::endl;
+      std::cout << "pre res: \n" << res.adj() << std::endl;
+  for (int k = 0; k < w.adj().outerSize(); ++k) {
+    for (typename T1::vari_type::InnerIterator
+             it(w.adj(), k);
          it; ++it) {
       it.valueRef() += res.adj().coeff(it.row()) * value_of(b).coeff(it.col());
     }
   }
+      std::cout << "post w adj: \n" << w.adj() << std::endl;
 }
 
 }  // namespace internal
@@ -79,7 +66,7 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
                                     const std::vector<int>& v,
                                     const std::vector<int>& u, const T2& b) {
   using sparse_val_mat
-      = Eigen::Map<const Eigen::SparseMatrix<double, Eigen::RowMajor>>;
+      = Eigen::Map<const Eigen::SparseMatrix<double>>;
   using sparse_dense_mul_type
       = decltype((std::declval<sparse_val_mat>() * value_of(b)).eval());
   using return_t = return_var_matrix_t<sparse_dense_mul_type, T1, T2>;
@@ -102,43 +89,36 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
                  [](auto&& x) { return x - 1; });
   if (!is_constant<T2>::value && !is_constant<T1>::value) {
     arena_t<promote_scalar_t<var, T2>> b_arena = b;
-    arena_t<promote_scalar_t<var, T1>> w_arena = to_arena(w);
-    auto w_val_arena = to_arena(value_of(w_arena));
-    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);
+    auto w_mat_arena = to_soa_sparse_matrix(m, n, w, u_arena, v_arena);
+    arena_t<return_t> res = w_mat_arena.val() * value_of(b_arena);
     reverse_pass_callback(
-        [m, n, w_arena, w_val_arena, v_arena, u_arena, res, b_arena]() mutable {
-          sparse_val_mat w_val_mat(m, n, w_val_arena.size(), u_arena.data(),
-                                   v_arena.data(), w_val_arena.data());
-          internal::update_w(w_arena, m, n, u_arena, v_arena, b_arena, res);
-          b_arena.adj() += w_val_mat.transpose() * res.adj();
+        [res, w_mat_arena, b_arena]() mutable {
+          internal::update_w(w_mat_arena, b_arena, res);
+          //w_mat_arena.adj()  += res.adj() * b_arena.val().transpose();
+          b_arena.adj() += w_mat_arena.val().transpose() * res.adj();
         });
+    std::cout << "NOT CALLED" << std::endl;
     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(
-        [m, n, w_val_arena, v_arena, u_arena, res, b_arena]() mutable {
-          sparse_val_mat w_val_mat(m, n, w_val_arena.size(), u_arena.data(),
-                                   v_arena.data(), w_val_arena.data());
+        [w_val_mat, res, b_arena]() mutable {
           b_arena.adj() += w_val_mat.transpose() * res.adj();
         });
+    std::cout << "NOT CALLED" << std::endl;
     return return_t(res);
   } else {
-    arena_t<promote_scalar_t<var, T1>> w_arena = to_arena(w);
-    auto&& w_val = eval(value_of(w_arena));
-    sparse_val_mat w_val_mat(m, n, w_val.size(), u_arena.data(), v_arena.data(),
-                             w_val.data());
+    auto w_mat_arena = to_soa_sparse_matrix(m, n, w, u_arena, v_arena);
     auto b_arena = to_arena(value_of(b));
-    arena_t<return_t> res = w_val_mat * b_arena;
+    arena_t<return_t> res = w_mat_arena.val() * b_arena;
     reverse_pass_callback(
-        [m, n, w_arena, v_arena, u_arena, res, b_arena]() mutable {
-          internal::update_w(w_arena, m, n, u_arena, v_arena, b_arena, res);
+        [res, w_mat_arena, b_arena]() mutable {
+          internal::update_w(w_mat_arena, b_arena, res);
+          std::cout << "w_mat: " << w_mat_arena.adj();
         });
     return return_t(res);
   }
diff --git a/stan/math/rev/fun/to_soa_sparse_matrix.hpp b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
new file mode 100644
index 00000000000..be95ad33476
--- /dev/null
+++ b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
@@ -0,0 +1,84 @@
+#ifndef STAN_MATH_REV_FUN_TO_SOA_SPARSE_MATRIX_HPP
+#define STAN_MATH_REV_FUN_TO_SOA_SPARSE_MATRIX_HPP
+
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/rev/core.hpp>
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/prim/meta.hpp>
+#include <type_traits>
+
+namespace stan {
+namespace math {
+
+template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
+  require_var_t<T1>* = nullptr,
+  require_eigen_dense_base_t<value_type_t<T1>>* = nullptr,
+  require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
+    arena_t<T2> v_arena(std::forward<T2>(v));
+    arena_t<T3> u_arena(std::forward<T3>(u));
+    using sparse_arena_mat_t = arena_t<Eigen::SparseMatrix<double, Options>>;
+    sparse_arena_mat_t arena_val_x(m, n, w.val().size(),
+      u_arena.data(), v_arena.data(), w.vi_->val_.data());
+    var_value<Eigen::SparseMatrix<double, Options>> var_x(arena_val_x);
+    sparse_arena_mat_t arena_adj_x(m, n, w.adj().size(),
+      u_arena.data(), v_arena.data(), w.vi_->adj_.data());
+    reverse_pass_callback([arena_adj_x, var_x]() mutable {
+      std::cout << "pre varx adj: \n" << var_x.adj() << std::endl;
+      std::cout << "pre arena_adj: \n" << arena_adj_x << std::endl;
+      arena_adj_x += var_x.adj();
+      std::cout << "post varx adj: \n" << var_x.adj() << std::endl;
+      std::cout << "post arena_adj: \n" << arena_adj_x << std::endl;
+    });
+    return var_x;
+  }
+
+
+template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
+  require_eigen_dense_base_vt<is_var, T1>* = nullptr,
+  require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
+    arena_t<T1> w_arena(std::forward<T1>(w));
+    arena_t<T2> v_arena(std::forward<T2>(v));
+    arena_t<T3> u_arena(std::forward<T3>(u));
+    arena_t<Eigen::SparseMatrix<var, Options>> arena_x(m, n, w_arena.size(),
+      u_arena.data(), v_arena.data(), w_arena.data());
+    var_value<Eigen::SparseMatrix<double, Options>> var_x(value_of(arena_x));
+    // No need to copy adj, but need to backprop
+    reverse_pass_callback([arena_x, var_x]() mutable {
+      using var_sparse_iterator_t = typename arena_t<Eigen::SparseMatrix<var, Options>>::InnerIterator;
+      using dbl_sparse_iterator_t = typename arena_t<Eigen::SparseMatrix<double, Options>>::InnerIterator;
+      // arena_x.adj() += var_x.adj() once custom adj() for var sparse matrix
+      for (int k = 0; k < arena_x.outerSize(); ++k) {
+        var_sparse_iterator_t it_arena_x(arena_x, k);
+        dbl_sparse_iterator_t it_var_x(var_x.adj(), k);
+        for (; bool(it_arena_x) && bool(it_var_x); ++it_arena_x, ++it_var_x) {
+          it_arena_x.valueRef().adj() += it_var_x.valueRef();
+        }
+      }
+    });
+    return var_x;
+  }
+
+template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
+  require_eigen_dense_base_vt<std::is_arithmetic, T1>* = nullptr,
+  require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
+  inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
+    arena_t<T1> w_arena(std::forward<T1>(w));
+    arena_t<T2> v_arena(std::forward<T2>(v));
+    arena_t<T3> u_arena(std::forward<T3>(u));
+    arena_t<Eigen::SparseMatrix<double, Options>> arena_x(m, n, w_arena.size(),
+      u_arena.data(), v_arena.data(), w_arena.data());
+    return var_value<Eigen::SparseMatrix<double, Options>>(arena_x);
+  }
+
+  template <int Options = Eigen::ColMajor, typename T, require_t<math::conjunction<is_var<T>, is_eigen_sparse_base<value_type_t<T>>>>* = nullptr>
+  inline auto to_soa_sparse_matrix(T&& x) {
+    return std::forward<T>(x);
+  }
+
+}
+}
+
+#endif
+
diff --git a/stan/math/rev/fun/to_var_value.hpp b/stan/math/rev/fun/to_var_value.hpp
index 004b520b6f2..6664b7d7521 100644
--- a/stan/math/rev/fun/to_var_value.hpp
+++ b/stan/math/rev/fun/to_var_value.hpp
@@ -18,7 +18,7 @@ namespace math {
  * @param a matrix to convert
  */
 template <typename T, require_eigen_vt<is_var, T>* = nullptr>
-var_value<Eigen::Matrix<double, T::RowsAtCompileTime, T::ColsAtCompileTime>>
+inline var_value<Eigen::Matrix<double, T::RowsAtCompileTime, T::ColsAtCompileTime>>
 to_var_value(const T& a) {
   arena_matrix<plain_type_t<T>> a_arena = a;
   var_value<promote_scalar_t<double, T>> res(a_arena.val());
@@ -34,7 +34,7 @@ to_var_value(const T& a) {
  * @param a matrix to convert
  */
 template <typename T, require_var_t<T>* = nullptr>
-T to_var_value(T&& a) {
+inline T to_var_value(T&& a) {
   return std::forward<T>(a);
 }
 
@@ -46,7 +46,7 @@ T to_var_value(T&& a) {
  * @param a std::vector of elements to convert
  */
 template <typename T>
-auto to_var_value(const std::vector<T>& a) {
+inline auto to_var_value(const std::vector<T>& a) {
   std::vector<decltype(to_var_value(std::declval<T>()))> out;
   out.reserve(a.size());
   for (size_t i = 0; i < a.size(); ++i) {
diff --git a/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp b/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
index 20c73955243..87e3377f417 100644
--- a/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
+++ b/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
@@ -1,7 +1,7 @@
 #include <test/unit/math/test_ad.hpp>
 #include <vector>
-
-TEST(MathMixMatFun, csr_matrix_times_vector) {
+/*
+TEST(MathMixMatFun, csr_matrix_times_vector1) {
   auto f = [](const auto& w, const auto& b) {
     using stan::math::csr_matrix_times_vector;
     std::vector<int> v{1, 2, 3, 1, 2};
@@ -15,5 +15,20 @@ TEST(MathMixMatFun, csr_matrix_times_vector) {
   b << 1, 2, 3, 4, 5;
 
   stan::test::expect_ad(f, w, b);
-  stan::test::expect_ad_matvar(f, w, b);
+}
+*/
+TEST(MathMixMatFun, csr_matrix_times_vector2) {
+  auto f = [](const auto& w) {
+    using stan::math::csr_matrix_times_vector;
+    std::vector<int> v{1, 2, 3, 1, 2};
+    std::vector<int> u{1, 2, 3, 4, 5, 6};
+  Eigen::VectorXd b(5);
+  b << 1, 2, 3, 4, 5;
+    return csr_matrix_times_vector(5, 5, w, v, u, b);
+  };
+
+  Eigen::VectorXd w(5);
+  w << -0.67082, 0.5, -0.223607, -0.223607, -0.5;
+
+  stan::test::expect_ad_matvar(f, w);
 }

From fec3689225e7d00172cabfcdcb7513ac6ecf959a Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 15 Apr 2024 15:25:53 -0400
Subject: [PATCH 2/9] update csr matrix multiply to avoid linker error for
 windows. Adds to_soa_sparse_matrix for making sparse matrices from csr
 formatted data

---
 stan/math/prim/meta/is_eigen_dense_base.hpp   | 16 ++++
 stan/math/rev/core/arena_matrix.hpp           | 25 +++---
 stan/math/rev/core/var.hpp                    | 12 +++
 stan/math/rev/core/vari.hpp                   | 29 ++++---
 stan/math/rev/fun/csr_matrix_times_vector.hpp | 70 +++++++---------
 stan/math/rev/fun/to_soa_sparse_matrix.hpp    | 79 +++++++++++++-----
 .../mix/fun/csr_matrix_times_vector_test.cpp  | 34 ++++++--
 .../rev/fun/to_soa_sparse_matrix_test.cpp     | 80 +++++++++++++++++++
 8 files changed, 259 insertions(+), 86 deletions(-)
 create mode 100644 test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp

diff --git a/stan/math/prim/meta/is_eigen_dense_base.hpp b/stan/math/prim/meta/is_eigen_dense_base.hpp
index 7b1ebc6cd1e..940f65bb2ef 100644
--- a/stan/math/prim/meta/is_eigen_dense_base.hpp
+++ b/stan/math/prim/meta/is_eigen_dense_base.hpp
@@ -33,6 +33,22 @@ using require_eigen_dense_base_t
     = require_t<is_eigen_dense_base<std::decay_t<T>>>;
 /*! @} */
 
+/*! \ingroup require_eigens_types */
+/*! \defgroup eigen_dense_base_types eigen_dense_base_types  */
+/*! \addtogroup eigen_dense_base_types */
+/*! @{ */
+
+/*! \brief Require type satisfies @ref is_eigen_dense_base */
+/*! and value type satisfies `TypeCheck` */
+/*! @tparam TypeCheck The type trait to check the value type against */
+/*! @tparam Check The type to test @ref is_eigen_dense_base for and whose
+ * @ref value_type is checked with `TypeCheck` */
+template <template <class...> class TypeCheck, class... Check>
+using require_eigen_dense_base_vt
+    = require_t<container_type_check_base<is_eigen_dense_base, value_type_t,
+                                          TypeCheck, Check...>>;
+/*! @} */
+
 }  // namespace stan
 
 #endif
diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 69baf5b81c0..769a09172f1 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -6,7 +6,6 @@
 #include <stan/math/rev/core/chainablestack.hpp>
 #include <stan/math/rev/core/var_value_fwd_declare.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
-
 namespace stan {
 namespace math {
 
@@ -225,8 +224,11 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   arena_matrix(const arena_matrix<MatrixType>& other)
       : Base::Map(other.rows(), other.cols(), other.nonZeros(),
-                  other.outerIndexPtr(), other.innerIndexPtr(),
-                  other.valuePtr(), other.innernonZeroPtr()) {}
+                  const_cast<StorageIndex*>(other.outerIndexPtr()),
+                  const_cast<StorageIndex*>(other.innerIndexPtr()),
+                  const_cast<Scalar*>(other.valuePtr()),
+                  const_cast<StorageIndex*>(other.innerNonZeroPtr())) {
+                  }
   /**
    * Move constructor.
    * @note Since the memory for the arena matrix sits in Stan's memory arena all
@@ -235,8 +237,10 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   arena_matrix(arena_matrix<MatrixType>&& other)
       : Base::Map(other.rows(), other.cols(), other.nonZeros(),
-                  other.outerIndexPtr(), other.innerIndexPtr(),
-                  other.valuePtr(), other.innerNonZeroPtr()) {}
+                  const_cast<StorageIndex*>(other.outerIndexPtr()),
+                  const_cast<StorageIndex*>(other.innerIndexPtr()),
+                  const_cast<Scalar*>(other.valuePtr()),
+                  const_cast<StorageIndex*>(other.innerNonZeroPtr())) {}
   /**
    * Copy constructor. No actual copy is performed
    * @note Since the memory for the arena matrix sits in Stan's memory arena all
@@ -245,8 +249,11 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   arena_matrix(arena_matrix<MatrixType>& other)
       : Base::Map(other.rows(), other.cols(), other.nonZeros(),
-                  other.outerIndexPtr(), other.innerIndexPtr(),
-                  other.valuePtr(), other.innerNonZeroPtr()) {}
+                    const_cast<StorageIndex*>(other.outerIndexPtr()),
+                    const_cast<StorageIndex*>(other.innerIndexPtr()),
+                    const_cast<Scalar*>(other.valuePtr()),
+                    const_cast<StorageIndex*>(other.innerNonZeroPtr())) {
+                    }
 
   // without this using, compiler prefers combination of implicit construction
   // and copy assignment to the inherited operator when assigned an expression
@@ -259,7 +266,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    * @return `*this`
    */
   template <typename ArenaMatrix,
-            require_same_t<ArenaMatrix, arena_matrix<MatrixType>>* = nullptr>
+            require_same_t<std::decay_t<ArenaMatrix>, arena_matrix<MatrixType>>* = nullptr>
   arena_matrix& operator=(ArenaMatrix&& other) {
     // placement new changes what data map points to - there is no allocation
     new (this) Base(other.rows(), other.cols(), other.nonZeros(),
@@ -280,7 +287,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
   template <typename Expr,
             require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
   arena_matrix& operator=(Expr&& expr) {
-    *this = arena_matrix(std::forward<Expr>(expr));
+    new (this) arena_matrix(std::forward<Expr>(expr));
     return *this;
   }
 
diff --git a/stan/math/rev/core/var.hpp b/stan/math/rev/core/var.hpp
index c56d4f09a42..c697af5361a 100644
--- a/stan/math/rev/core/var.hpp
+++ b/stan/math/rev/core/var.hpp
@@ -418,6 +418,18 @@ class var_value<T, internal::require_matrix_var_value<T>> {
         });
   }
 
+  /**
+   * Construct a `var_value` with a plain type
+   *  from another `var_value` containing an expression.
+   * @tparam S type of the value in the `var_value` to assing
+   * @param other the value to assign
+   * @return this
+   */
+  template <typename S, typename T_ = T,
+            require_assignable_t<value_type, S>* = nullptr,
+            require_arena_matrix_t<S>* = nullptr>
+  var_value(const S& val, const S& adj) : vi_(new vari_type(val, adj)) {}
+
   /**
    * Construct a variable from a pointer to a variable implementation.
    * @param vi A vari_value pointer.
diff --git a/stan/math/rev/core/vari.hpp b/stan/math/rev/core/vari.hpp
index 80da3b0f46d..9d34447f997 100644
--- a/stan/math/rev/core/vari.hpp
+++ b/stan/math/rev/core/vari.hpp
@@ -821,17 +821,16 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    */
   static constexpr int ColsAtCompileTime = T::ColsAtCompileTime;
 
+  /**
+   * The value of this variable.
+   */
+  arena_matrix<PlainObject> val_;
   /**
    * The adjoint of this variable, which is the partial derivative
    * of this variable with respect to the root variable.
    */
   arena_matrix<PlainObject> adj_;
 
-  /**
-   * The value of this variable.
-   */
-  arena_matrix<PlainObject> val_;
-
   /**
    * Construct a variable implementation from a value. The
    * adjoint is initialized to zero.
@@ -847,10 +846,19 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    * @param x Value of the constructed variable.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  explicit vari_value(S&& x) : adj_(x), val_(std::forward<S>(x)) {
-    this->set_zero_adjoint();
+  explicit vari_value(S&& x) : val_(std::forward<S>(x)),
+    adj_(val_.rows(), val_.cols(), val_.nonZeros(),
+      val_.outerIndexPtr(), val_.innerIndexPtr(),
+      arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
+      val_.innerNonZeroPtr()) {
+    ChainableStack::instance_->var_stack_.push_back(this);
+  }
+
+  vari_value(const arena_matrix<PlainObject>& val,
+   const arena_matrix<PlainObject>& adj) : val_(val), adj_(adj) {
     ChainableStack::instance_->var_stack_.push_back(this);
   }
+
   /**
    * Construct an sparse Eigen variable implementation from a value. The
    *  adjoint is initialized to zero and if `stacked` is `false` this vari
@@ -869,8 +877,11 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    * that its `chain()` method is not called.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  vari_value(S&& x, bool stacked) : adj_(x), val_(std::forward<S>(x)) {
-    this->set_zero_adjoint();
+  vari_value(S&& x, bool stacked) : val_(std::forward<S>(x)),
+    adj_(val_.rows(), val_.cols(), val_.nonZeros(),
+      val_.outerIndexPtr(), val_.innerIndexPtr(),
+      arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
+      val_.innerNonZeroPtr()) {
     if (stacked) {
       ChainableStack::instance_->var_stack_.push_back(this);
     } else {
diff --git a/stan/math/rev/fun/csr_matrix_times_vector.hpp b/stan/math/rev/fun/csr_matrix_times_vector.hpp
index 8857f1ce23d..9af92b8d83a 100644
--- a/stan/math/rev/fun/csr_matrix_times_vector.hpp
+++ b/stan/math/rev/fun/csr_matrix_times_vector.hpp
@@ -11,26 +11,6 @@
 namespace stan {
 namespace math {
 
-namespace internal {
-
-template <typename T1, typename T2, typename Res,
-          require_var_matrix_t<T1>* = nullptr>
-inline void update_w(T1& w, T2&& b, Res&& res) {
-      std::cout << "pre w adj: \n" << w.adj() << std::endl;
-      std::cout << "pre b: \n" << value_of(b).eval() << std::endl;
-      std::cout << "pre res: \n" << res.adj() << std::endl;
-  for (int k = 0; k < w.adj().outerSize(); ++k) {
-    for (typename T1::vari_type::InnerIterator
-             it(w.adj(), k);
-         it; ++it) {
-      it.valueRef() += res.adj().coeff(it.row()) * value_of(b).coeff(it.col());
-    }
-  }
-      std::cout << "post w adj: \n" << w.adj() << std::endl;
-}
-
-}  // namespace internal
-
 /**
  * \addtogroup csr_format
  * Return the multiplication of the sparse matrix (specified by
@@ -66,7 +46,7 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
                                     const std::vector<int>& v,
                                     const std::vector<int>& u, const T2& b) {
   using sparse_val_mat
-      = Eigen::Map<const Eigen::SparseMatrix<double>>;
+      = Eigen::Map<const Eigen::SparseMatrix<double, Eigen::RowMajor>>;
   using sparse_dense_mul_type
       = decltype((std::declval<sparse_val_mat>() * value_of(b)).eval());
   using return_t = return_var_matrix_t<sparse_dense_mul_type, T1, T2>;
@@ -87,17 +67,24 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
   std::vector<int, arena_allocator<int>> u_arena(u.size());
   std::transform(u.begin(), u.end(), u_arena.begin(),
                  [](auto&& x) { return x - 1; });
+  using sparse_var_value_t
+      = var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>>;
   if (!is_constant<T2>::value && !is_constant<T1>::value) {
     arena_t<promote_scalar_t<var, T2>> b_arena = b;
-    auto w_mat_arena = to_soa_sparse_matrix(m, n, w, u_arena, v_arena);
+    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 {
-          internal::update_w(w_mat_arena, b_arena, res);
-          //w_mat_arena.adj()  += res.adj() * b_arena.val().transpose();
-          b_arena.adj() += w_mat_arena.val().transpose() * res.adj();
-        });
-    std::cout << "NOT CALLED" << std::endl;
+    reverse_pass_callback([res, w_mat_arena, b_arena]() mutable {
+      for (int k = 0; k < w_mat_arena.adj().outerSize(); ++k) {
+        for (typename sparse_var_value_t::vari_type::InnerIterator it(
+                 w_mat_arena.adj(), k);
+             it; ++it) {
+          it.valueRef()
+              += res.adj().coeff(it.row()) * b_arena.val().coeff(it.col());
+        }
+      }
+      b_arena.adj() += w_mat_arena.val().transpose() * res.adj();
+    });
     return return_t(res);
   } else if (!is_constant<T2>::value) {
     arena_t<promote_scalar_t<var, T2>> b_arena = b;
@@ -105,21 +92,24 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& 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();
-        });
-    std::cout << "NOT CALLED" << std::endl;
+    reverse_pass_callback([w_val_mat, res, b_arena]() mutable {
+      b_arena.adj() += w_val_mat.transpose() * res.adj();
+    });
     return return_t(res);
   } else {
-    auto w_mat_arena = to_soa_sparse_matrix(m, n, w, u_arena, v_arena);
+    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 {
-          internal::update_w(w_mat_arena, b_arena, res);
-          std::cout << "w_mat: " << w_mat_arena.adj();
-        });
+    reverse_pass_callback([res, w_mat_arena, b_arena]() mutable {
+      for (int k = 0; k < w_mat_arena.adj().outerSize(); ++k) {
+        for (typename sparse_var_value_t::vari_type::InnerIterator it(
+                 w_mat_arena.adj(), k);
+             it; ++it) {
+          it.valueRef() += res.adj().coeff(it.row()) * b_arena.coeff(it.col());
+        }
+      }
+    });
     return return_t(res);
   }
 }
diff --git a/stan/math/rev/fun/to_soa_sparse_matrix.hpp b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
index be95ad33476..5612bf0da22 100644
--- a/stan/math/rev/fun/to_soa_sparse_matrix.hpp
+++ b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
@@ -10,37 +10,62 @@
 namespace stan {
 namespace math {
 
+
+/**
+ * Create a sparse matrix from the given SoA matrix and indexes.
+ * @tparam Options Eigen matrix options.
+ * @tparam T1 A @ref var_value with a dense vector inner type
+ * @tparam T2 Container type of the column indexes.
+ * @tparam T3 Container type of the row indexes.
+ * @param m Number of rows in matrix.
+ * @param n Number of columns in matrix.
+ * @param w Vector of non-zero values in matrix.
+ * @param u Index of where each row starts in w, length equal to
+ *          the number of rows plus one.
+ * @param v Column index of each non-zero value, same
+ *          length as w.
+ * @return Sparse matrix.
+ */
 template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
   require_var_t<T1>* = nullptr,
   require_eigen_dense_base_t<value_type_t<T1>>* = nullptr,
   require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
 inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
-    arena_t<T2> v_arena(std::forward<T2>(v));
-    arena_t<T3> u_arena(std::forward<T3>(u));
-    using sparse_arena_mat_t = arena_t<Eigen::SparseMatrix<double, Options>>;
+    auto v_arena = to_arena(v);
+    auto u_arena = to_arena(u);
+    using sparse_mat_t = Eigen::SparseMatrix<double, Options>;
+    using sparse_arena_mat_t = arena_t<sparse_mat_t>;
     sparse_arena_mat_t arena_val_x(m, n, w.val().size(),
       u_arena.data(), v_arena.data(), w.vi_->val_.data());
-    var_value<Eigen::SparseMatrix<double, Options>> var_x(arena_val_x);
     sparse_arena_mat_t arena_adj_x(m, n, w.adj().size(),
       u_arena.data(), v_arena.data(), w.vi_->adj_.data());
-    reverse_pass_callback([arena_adj_x, var_x]() mutable {
-      std::cout << "pre varx adj: \n" << var_x.adj() << std::endl;
-      std::cout << "pre arena_adj: \n" << arena_adj_x << std::endl;
-      arena_adj_x += var_x.adj();
-      std::cout << "post varx adj: \n" << var_x.adj() << std::endl;
-      std::cout << "post arena_adj: \n" << arena_adj_x << std::endl;
-    });
+    var_value<sparse_mat_t> var_x(arena_val_x, arena_adj_x);
     return var_x;
   }
 
 
+/**
+ * Create a sparse matrix from the given AoS matrix of vars and indexes.
+ * @tparam Options Eigen matrix options.
+ * @tparam T1 A type inheriting from `Eigen::DenseBase` with a scalar type of @ref var_value
+ * @tparam T2 Container type of the column indexes.
+ * @tparam T3 Container type of the row indexes.
+ * @param m Number of rows in matrix.
+ * @param n Number of columns in matrix.
+ * @param w Vector of non-zero values in matrix.
+ * @param u Index of where each row starts in w, length equal to
+ *          the number of rows plus one.
+ * @param v Column index of each non-zero value, same
+ *          length as w.
+ * @return Sparse matrix.
+ */
 template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
   require_eigen_dense_base_vt<is_var, T1>* = nullptr,
   require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
 inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
-    arena_t<T1> w_arena(std::forward<T1>(w));
-    arena_t<T2> v_arena(std::forward<T2>(v));
-    arena_t<T3> u_arena(std::forward<T3>(u));
+    auto w_arena = to_arena(w);
+    auto v_arena = to_arena(v);
+    auto u_arena = to_arena(u);
     arena_t<Eigen::SparseMatrix<var, Options>> arena_x(m, n, w_arena.size(),
       u_arena.data(), v_arena.data(), w_arena.data());
     var_value<Eigen::SparseMatrix<double, Options>> var_x(value_of(arena_x));
@@ -60,23 +85,33 @@ inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
     return var_x;
   }
 
+/**
+ * Create a sparse matrix from the given matrix of floats and indexes.
+ * @tparam Options Eigen matrix options.
+ * @tparam T1 A type inheriting from `Eigen::DenseBase` with an arithmetic scalar type
+ * @tparam T2 Container type of the column indexes.
+ * @tparam T3 Container type of the row indexes.
+ * @param m Number of rows in matrix.
+ * @param n Number of columns in matrix.
+ * @param w Vector of non-zero values in matrix.
+ * @param u Index of where each row starts in w, length equal to
+ *          the number of rows plus one.
+ * @param v Column index of each non-zero value, same
+ *          length as w.
+ * @return Sparse matrix.
+ */
 template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
   require_eigen_dense_base_vt<std::is_arithmetic, T1>* = nullptr,
   require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
   inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
-    arena_t<T1> w_arena(std::forward<T1>(w));
-    arena_t<T2> v_arena(std::forward<T2>(v));
-    arena_t<T3> u_arena(std::forward<T3>(u));
+    auto w_arena = to_arena(w);
+    auto v_arena = to_arena(v);
+    auto u_arena = to_arena(u);
     arena_t<Eigen::SparseMatrix<double, Options>> arena_x(m, n, w_arena.size(),
       u_arena.data(), v_arena.data(), w_arena.data());
     return var_value<Eigen::SparseMatrix<double, Options>>(arena_x);
   }
 
-  template <int Options = Eigen::ColMajor, typename T, require_t<math::conjunction<is_var<T>, is_eigen_sparse_base<value_type_t<T>>>>* = nullptr>
-  inline auto to_soa_sparse_matrix(T&& x) {
-    return std::forward<T>(x);
-  }
-
 }
 }
 
diff --git a/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp b/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
index 87e3377f417..f17980e3d29 100644
--- a/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
+++ b/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
@@ -1,6 +1,27 @@
 #include <test/unit/math/test_ad.hpp>
 #include <vector>
-/*
+
+TEST(MathMixMatFun, csr_matrix_times_vector_vals) {
+  using stan::math::csr_matrix_times_vector;
+  std::vector<int> v{1, 2, 3, 1, 2};
+  std::vector<int> u{1, 2, 3, 4, 5, 6};
+  Eigen::VectorXd b(5);
+  b << 1, 2, 3, 4, 5;
+  Eigen::VectorXd w(5);
+  w << -0.67082, 0.5, -0.223607, -0.223607, -0.5;
+  auto dbl_res = csr_matrix_times_vector(5, 5, w, v, u, b);
+  using stan::math::var;
+  Eigen::Matrix<var, -1, 1> b_var(b);
+  Eigen::Matrix<var, -1, 1> w_var(w);
+  auto var_res_w = csr_matrix_times_vector(5, 5, w_var, v, u, b);
+  EXPECT_MATRIX_EQ(dbl_res, var_res_w);
+  auto var_res_b = csr_matrix_times_vector(5, 5, w, v, u, b_var);
+  EXPECT_MATRIX_EQ(dbl_res, var_res_b);
+  auto var_res_bw = csr_matrix_times_vector(5, 5, w_var, v, u, b_var);
+  EXPECT_MATRIX_EQ(dbl_res, var_res_bw);
+}
+
+
 TEST(MathMixMatFun, csr_matrix_times_vector1) {
   auto f = [](const auto& w, const auto& b) {
     using stan::math::csr_matrix_times_vector;
@@ -16,19 +37,20 @@ TEST(MathMixMatFun, csr_matrix_times_vector1) {
 
   stan::test::expect_ad(f, w, b);
 }
-*/
+
 TEST(MathMixMatFun, csr_matrix_times_vector2) {
-  auto f = [](const auto& w) {
+  auto f = [](const auto& w, const auto& b) {
     using stan::math::csr_matrix_times_vector;
     std::vector<int> v{1, 2, 3, 1, 2};
     std::vector<int> u{1, 2, 3, 4, 5, 6};
-  Eigen::VectorXd b(5);
-  b << 1, 2, 3, 4, 5;
     return csr_matrix_times_vector(5, 5, w, v, u, b);
   };
 
   Eigen::VectorXd w(5);
   w << -0.67082, 0.5, -0.223607, -0.223607, -0.5;
+  Eigen::VectorXd b(5);
+  b << 1, 2, 3, 4, 5;
 
-  stan::test::expect_ad_matvar(f, w);
+  stan::test::expect_ad_matvar(f, w, b);
 }
+
diff --git a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
new file mode 100644
index 00000000000..70db6bc3571
--- /dev/null
+++ b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
@@ -0,0 +1,80 @@
+#include <iostream>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/rev/fun/to_soa_sparse_matrix.hpp>
+#include <stan/math/rev/core.hpp>
+#include <test/unit/math/rev/util.hpp>
+#include <vector>
+/*
+TEST_F(AgradRev, to_soa_sparse_matrix_matrix_double) {
+  using stan::math::var;
+  using stan::math::var_value;
+  using stan::math::to_soa_sparse_matrix;
+  std::vector<int> v{0, 1, 2, 0, 1};
+  std::vector<int> u{0, 1, 2, 3, 4, 5};
+  Eigen::VectorXd w(5);
+  int m = 5;
+  int n = 5;
+  w << 1, 2, 3, 4, 5;
+  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena =
+    to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w, u, v);
+  for (int i = 0; i < 5; ++i) {
+    EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w(i));
+  }
+}
+
+
+TEST_F(AgradRev, to_soa_sparse_matrix_matrix_var) {
+  using stan::math::var;
+  using stan::math::var_value;
+  using stan::math::to_soa_sparse_matrix;
+  std::vector<int> v{0, 1, 2, 0, 1};
+  std::vector<int> u{0, 1, 2, 3, 4, 5};
+  Eigen::VectorXd w(5);
+  int m = 5;
+  int n = 5;
+  w << 1, 2, 3, 4, 5;
+  Eigen::Matrix<var, -1, 1> w_var(w);
+  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena =
+    to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w_var, u, v);
+  for (int i = 0; i < 5; ++i) {
+    EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w.val()(i));
+  }
+}
+*/
+
+TEST_F(AgradRev, to_soa_sparse_matrix_var_matrix) {
+  using stan::math::var;
+  using stan::math::var_value;
+  using stan::math::to_soa_sparse_matrix;
+  std::vector<int> v{0, 1, 2, 0, 1};
+  std::vector<int> u{0, 1, 2, 3, 4, 5};
+  Eigen::VectorXd w(5);
+  int m = 5;
+  int n = 5;
+  w << 1, 2, 3, 4, 5;
+  var_value<Eigen::VectorXd> w_var(w);
+  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena =
+    to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w_var, u, v);
+  for (int i = 0; i < 5; ++i) {
+    EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w.val()(i));
+  }
+  // Changing this value should change the adjoint of the sparse matrix
+  for (int i = 0; i < 5; ++i) {
+    w_mat_arena.adj().valuePtr()[i] = i;
+  }
+  for (int i = 0; i < 5; ++i) {
+    EXPECT_EQ(w_mat_arena.adj().valuePtr()[i], w_var.adj()(i));
+  }
+  for (int i = 0; i < 5; ++i) {
+    w_mat_arena.adj().valuePtr()[i] = 0;
+  }
+  for (int i = 0; i < 5; ++i) {
+    w_var.adj().coeffRef(i) = i;
+  }
+  for (int i = 0; i < 5; ++i) {
+    EXPECT_EQ(w_mat_arena.adj().valuePtr()[i], w_var.adj()(i));
+  }
+
+ }
+
+

From 6b8ae15433644303449529e81b42f695d3b0ea7b Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 15 Apr 2024 15:54:22 -0400
Subject: [PATCH 3/9] small fixes

---
 stan/math/prim/fun/value_of.hpp            |  1 -
 stan/math/rev/core/var.hpp                 |  8 +--
 stan/math/rev/fun/to_soa_sparse_matrix.hpp | 60 +++++++++++-----------
 3 files changed, 34 insertions(+), 35 deletions(-)

diff --git a/stan/math/prim/fun/value_of.hpp b/stan/math/prim/fun/value_of.hpp
index cdb5115d104..a4ff276ac54 100644
--- a/stan/math/prim/fun/value_of.hpp
+++ b/stan/math/prim/fun/value_of.hpp
@@ -1,7 +1,6 @@
 #ifndef STAN_MATH_PRIM_FUN_VALUE_OF_HPP
 #define STAN_MATH_PRIM_FUN_VALUE_OF_HPP
 
-#include <iostream>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <cstddef>
diff --git a/stan/math/rev/core/var.hpp b/stan/math/rev/core/var.hpp
index c697af5361a..bce9283225d 100644
--- a/stan/math/rev/core/var.hpp
+++ b/stan/math/rev/core/var.hpp
@@ -419,11 +419,11 @@ class var_value<T, internal::require_matrix_var_value<T>> {
   }
 
   /**
-   * Construct a `var_value` with a plain type
-   *  from another `var_value` containing an expression.
+   * Construct a `var_value` with premade @ref arena_matrix types.
+   *  The values and adjoint matrices passed here will be shallow copied.
    * @tparam S type of the value in the `var_value` to assing
-   * @param other the value to assign
-   * @return this
+   * @param val The value matrix to go into the vari
+   * @param adj the adjoint matrix to go into the vari
    */
   template <typename S, typename T_ = T,
             require_assignable_t<value_type, S>* = nullptr,
diff --git a/stan/math/rev/fun/to_soa_sparse_matrix.hpp b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
index 5612bf0da22..038076c0987 100644
--- a/stan/math/rev/fun/to_soa_sparse_matrix.hpp
+++ b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
@@ -14,9 +14,9 @@ namespace math {
 /**
  * Create a sparse matrix from the given SoA matrix and indexes.
  * @tparam Options Eigen matrix options.
- * @tparam T1 A @ref var_value with a dense vector inner type
- * @tparam T2 Container type of the column indexes.
- * @tparam T3 Container type of the row indexes.
+ * @tparam VarMatrix A @ref var_value with a dense vector inner type
+ * @tparam Vec1 Container type of the column indexes.
+ * @tparam Vec2 Container type of the row indexes.
  * @param m Number of rows in matrix.
  * @param n Number of columns in matrix.
  * @param w Vector of non-zero values in matrix.
@@ -26,13 +26,13 @@ namespace math {
  *          length as w.
  * @return Sparse matrix.
  */
-template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
-  require_var_t<T1>* = nullptr,
-  require_eigen_dense_base_t<value_type_t<T1>>* = nullptr,
-  require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
-inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
-    auto v_arena = to_arena(v);
-    auto u_arena = to_arena(u);
+template <int Options = Eigen::ColMajor, typename VarMatrix, typename Vec1, typename Vec2,
+  require_var_t<VarMatrix>* = nullptr,
+  require_eigen_dense_base_t<value_type_t<VarMatrix>>* = nullptr,
+  require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, VarMatrix&& w, Vec1&& u, Vec2&& v) {
+    auto u_arena = to_arena(std::forward<Vec1>(u));
+    auto v_arena = to_arena(std::forward<Vec2>(v));
     using sparse_mat_t = Eigen::SparseMatrix<double, Options>;
     using sparse_arena_mat_t = arena_t<sparse_mat_t>;
     sparse_arena_mat_t arena_val_x(m, n, w.val().size(),
@@ -47,9 +47,9 @@ inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
 /**
  * Create a sparse matrix from the given AoS matrix of vars and indexes.
  * @tparam Options Eigen matrix options.
- * @tparam T1 A type inheriting from `Eigen::DenseBase` with a scalar type of @ref var_value
- * @tparam T2 Container type of the column indexes.
- * @tparam T3 Container type of the row indexes.
+ * @tparam MatrixVar A type inheriting from `Eigen::DenseBase` with a scalar type of @ref var_value
+ * @tparam Vec1 Container type of the column indexes.
+ * @tparam Vec2 Container type of the row indexes.
  * @param m Number of rows in matrix.
  * @param n Number of columns in matrix.
  * @param w Vector of non-zero values in matrix.
@@ -59,13 +59,13 @@ inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
  *          length as w.
  * @return Sparse matrix.
  */
-template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
-  require_eigen_dense_base_vt<is_var, T1>* = nullptr,
-  require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
-inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
-    auto w_arena = to_arena(w);
-    auto v_arena = to_arena(v);
-    auto u_arena = to_arena(u);
+template <int Options = Eigen::ColMajor, typename MatrixVar, typename Vec1, typename Vec2,
+  require_eigen_dense_base_vt<is_var, MatrixVar>* = nullptr,
+  require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, MatrixVar&& w, Vec1&& u, Vec2&& v) {
+    auto w_arena = to_arena(std::forward<MatrixVar>(w));
+    auto u_arena = to_arena(std::forward<Vec1>(u));
+    auto v_arena = to_arena(std::forward<Vec2>(v));
     arena_t<Eigen::SparseMatrix<var, Options>> arena_x(m, n, w_arena.size(),
       u_arena.data(), v_arena.data(), w_arena.data());
     var_value<Eigen::SparseMatrix<double, Options>> var_x(value_of(arena_x));
@@ -88,9 +88,9 @@ inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
 /**
  * Create a sparse matrix from the given matrix of floats and indexes.
  * @tparam Options Eigen matrix options.
- * @tparam T1 A type inheriting from `Eigen::DenseBase` with an arithmetic scalar type
- * @tparam T2 Container type of the column indexes.
- * @tparam T3 Container type of the row indexes.
+ * @tparam Mat A type inheriting from `Eigen::DenseBase` with an arithmetic scalar type
+ * @tparam Vec1 Container type of the column indexes.
+ * @tparam Vec2 Container type of the row indexes.
  * @param m Number of rows in matrix.
  * @param n Number of columns in matrix.
  * @param w Vector of non-zero values in matrix.
@@ -100,13 +100,13 @@ inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
  *          length as w.
  * @return Sparse matrix.
  */
-template <int Options = Eigen::ColMajor, typename T1, typename T2, typename T3,
-  require_eigen_dense_base_vt<std::is_arithmetic, T1>* = nullptr,
-  require_all_std_vector_vt<std::is_integral, T2, T3>* = nullptr>
-  inline auto to_soa_sparse_matrix(int m, int n, T1&& w, T2&& u, T3&& v) {
-    auto w_arena = to_arena(w);
-    auto v_arena = to_arena(v);
-    auto u_arena = to_arena(u);
+template <int Options = Eigen::ColMajor, typename Mat, typename Vec1, typename Vec2,
+  require_eigen_dense_base_vt<std::is_arithmetic, Mat>* = nullptr,
+  require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
+  inline auto to_soa_sparse_matrix(int m, int n, Mat&& w, Vec1&& u, Vec2&& v) {
+    auto w_arena = to_arena(std::forward<Mat>(w));
+    auto u_arena = to_arena(std::forward<Vec1>(u));
+    auto v_arena = to_arena(std::forward<Vec2>(v));
     arena_t<Eigen::SparseMatrix<double, Options>> arena_x(m, n, w_arena.size(),
       u_arena.data(), v_arena.data(), w_arena.data());
     return var_value<Eigen::SparseMatrix<double, Options>>(arena_x);

From c83cbfcc9f63a42b8521a5e6d57b14d1edcc0f6e Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Mon, 15 Apr 2024 16:13:44 -0400
Subject: [PATCH 4/9] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 stan/math/prim/fun/value_of.hpp               |   3 +-
 stan/math/prim/meta/promote_scalar_type.hpp   |   5 +-
 stan/math/rev/core/arena_matrix.hpp           |  15 +--
 stan/math/rev/core/vari.hpp                   |  25 ++--
 stan/math/rev/fun/to_soa_sparse_matrix.hpp    | 121 +++++++++---------
 stan/math/rev/fun/to_var_value.hpp            |   3 +-
 .../mix/fun/csr_matrix_times_vector_test.cpp  |   2 -
 .../rev/fun/to_soa_sparse_matrix_test.cpp     |  11 +-
 8 files changed, 95 insertions(+), 90 deletions(-)

diff --git a/stan/math/prim/fun/value_of.hpp b/stan/math/prim/fun/value_of.hpp
index a4ff276ac54..64286bbe0c3 100644
--- a/stan/math/prim/fun/value_of.hpp
+++ b/stan/math/prim/fun/value_of.hpp
@@ -82,7 +82,8 @@ template <typename EigMat, require_eigen_sparse_base_t<EigMat>* = nullptr,
 inline auto value_of(EigMat&& M) {
   auto&& M_ref = to_ref(M);
   using scalar_t = decltype(value_of(std::declval<value_type_t<EigMat>>()));
-  promote_scalar_t<scalar_t, plain_type_t<EigMat>> ret(M_ref.rows(), M_ref.cols());
+  promote_scalar_t<scalar_t, plain_type_t<EigMat>> ret(M_ref.rows(),
+                                                       M_ref.cols());
   ret.reserve(M_ref.nonZeros());
   for (int k = 0; k < M_ref.outerSize(); ++k) {
     for (typename std::decay_t<EigMat>::InnerIterator it(M_ref, k); it; ++it) {
diff --git a/stan/math/prim/meta/promote_scalar_type.hpp b/stan/math/prim/meta/promote_scalar_type.hpp
index 0c30b5a6140..da1a2deb405 100644
--- a/stan/math/prim/meta/promote_scalar_type.hpp
+++ b/stan/math/prim/meta/promote_scalar_type.hpp
@@ -98,8 +98,9 @@ struct promote_scalar_type<T, S, require_eigen_sparse_base_t<S>> {
   /**
    * The promoted type.
    */
-  using type = Eigen::SparseMatrix<typename promote_scalar_type<T, typename S::Scalar>::type,
-   S::Options, typename S::StorageIndex>;
+  using type = Eigen::SparseMatrix<
+      typename promote_scalar_type<T, typename S::Scalar>::type, S::Options,
+      typename S::StorageIndex>;
 };
 
 template <typename... PromotionScalars, typename... UnPromotedTypes>
diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 769a09172f1..e76f5af8429 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -227,8 +227,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
                   const_cast<StorageIndex*>(other.outerIndexPtr()),
                   const_cast<StorageIndex*>(other.innerIndexPtr()),
                   const_cast<Scalar*>(other.valuePtr()),
-                  const_cast<StorageIndex*>(other.innerNonZeroPtr())) {
-                  }
+                  const_cast<StorageIndex*>(other.innerNonZeroPtr())) {}
   /**
    * Move constructor.
    * @note Since the memory for the arena matrix sits in Stan's memory arena all
@@ -249,11 +248,10 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   arena_matrix(arena_matrix<MatrixType>& other)
       : Base::Map(other.rows(), other.cols(), other.nonZeros(),
-                    const_cast<StorageIndex*>(other.outerIndexPtr()),
-                    const_cast<StorageIndex*>(other.innerIndexPtr()),
-                    const_cast<Scalar*>(other.valuePtr()),
-                    const_cast<StorageIndex*>(other.innerNonZeroPtr())) {
-                    }
+                  const_cast<StorageIndex*>(other.outerIndexPtr()),
+                  const_cast<StorageIndex*>(other.innerIndexPtr()),
+                  const_cast<Scalar*>(other.valuePtr()),
+                  const_cast<StorageIndex*>(other.innerNonZeroPtr())) {}
 
   // without this using, compiler prefers combination of implicit construction
   // and copy assignment to the inherited operator when assigned an expression
@@ -266,7 +264,8 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    * @return `*this`
    */
   template <typename ArenaMatrix,
-            require_same_t<std::decay_t<ArenaMatrix>, arena_matrix<MatrixType>>* = nullptr>
+            require_same_t<std::decay_t<ArenaMatrix>,
+                           arena_matrix<MatrixType>>* = nullptr>
   arena_matrix& operator=(ArenaMatrix&& other) {
     // placement new changes what data map points to - there is no allocation
     new (this) Base(other.rows(), other.cols(), other.nonZeros(),
diff --git a/stan/math/rev/core/vari.hpp b/stan/math/rev/core/vari.hpp
index 9d34447f997..cef37d29a31 100644
--- a/stan/math/rev/core/vari.hpp
+++ b/stan/math/rev/core/vari.hpp
@@ -846,16 +846,18 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    * @param x Value of the constructed variable.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  explicit vari_value(S&& x) : val_(std::forward<S>(x)),
-    adj_(val_.rows(), val_.cols(), val_.nonZeros(),
-      val_.outerIndexPtr(), val_.innerIndexPtr(),
-      arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
-      val_.innerNonZeroPtr()) {
+  explicit vari_value(S&& x)
+      : val_(std::forward<S>(x)),
+        adj_(val_.rows(), val_.cols(), val_.nonZeros(), val_.outerIndexPtr(),
+             val_.innerIndexPtr(),
+             arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
+             val_.innerNonZeroPtr()) {
     ChainableStack::instance_->var_stack_.push_back(this);
   }
 
   vari_value(const arena_matrix<PlainObject>& val,
-   const arena_matrix<PlainObject>& adj) : val_(val), adj_(adj) {
+             const arena_matrix<PlainObject>& adj)
+      : val_(val), adj_(adj) {
     ChainableStack::instance_->var_stack_.push_back(this);
   }
 
@@ -877,11 +879,12 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    * that its `chain()` method is not called.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  vari_value(S&& x, bool stacked) : val_(std::forward<S>(x)),
-    adj_(val_.rows(), val_.cols(), val_.nonZeros(),
-      val_.outerIndexPtr(), val_.innerIndexPtr(),
-      arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
-      val_.innerNonZeroPtr()) {
+  vari_value(S&& x, bool stacked)
+      : val_(std::forward<S>(x)),
+        adj_(val_.rows(), val_.cols(), val_.nonZeros(), val_.outerIndexPtr(),
+             val_.innerIndexPtr(),
+             arena_matrix<Eigen::VectorXd>(val_.nonZeros()).setZero().data(),
+             val_.innerNonZeroPtr()) {
     if (stacked) {
       ChainableStack::instance_->var_stack_.push_back(this);
     } else {
diff --git a/stan/math/rev/fun/to_soa_sparse_matrix.hpp b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
index 038076c0987..2b87a4ea165 100644
--- a/stan/math/rev/fun/to_soa_sparse_matrix.hpp
+++ b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
@@ -10,7 +10,6 @@
 namespace stan {
 namespace math {
 
-
 /**
  * Create a sparse matrix from the given SoA matrix and indexes.
  * @tparam Options Eigen matrix options.
@@ -26,28 +25,29 @@ namespace math {
  *          length as w.
  * @return Sparse matrix.
  */
-template <int Options = Eigen::ColMajor, typename VarMatrix, typename Vec1, typename Vec2,
-  require_var_t<VarMatrix>* = nullptr,
-  require_eigen_dense_base_t<value_type_t<VarMatrix>>* = nullptr,
-  require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
-inline auto to_soa_sparse_matrix(int m, int n, VarMatrix&& w, Vec1&& u, Vec2&& v) {
-    auto u_arena = to_arena(std::forward<Vec1>(u));
-    auto v_arena = to_arena(std::forward<Vec2>(v));
-    using sparse_mat_t = Eigen::SparseMatrix<double, Options>;
-    using sparse_arena_mat_t = arena_t<sparse_mat_t>;
-    sparse_arena_mat_t arena_val_x(m, n, w.val().size(),
-      u_arena.data(), v_arena.data(), w.vi_->val_.data());
-    sparse_arena_mat_t arena_adj_x(m, n, w.adj().size(),
-      u_arena.data(), v_arena.data(), w.vi_->adj_.data());
-    var_value<sparse_mat_t> var_x(arena_val_x, arena_adj_x);
-    return var_x;
-  }
-
+template <int Options = Eigen::ColMajor, typename VarMatrix, typename Vec1,
+          typename Vec2, require_var_t<VarMatrix>* = nullptr,
+          require_eigen_dense_base_t<value_type_t<VarMatrix>>* = nullptr,
+          require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, VarMatrix&& w, Vec1&& u,
+                                 Vec2&& v) {
+  auto u_arena = to_arena(std::forward<Vec1>(u));
+  auto v_arena = to_arena(std::forward<Vec2>(v));
+  using sparse_mat_t = Eigen::SparseMatrix<double, Options>;
+  using sparse_arena_mat_t = arena_t<sparse_mat_t>;
+  sparse_arena_mat_t arena_val_x(m, n, w.val().size(), u_arena.data(),
+                                 v_arena.data(), w.vi_->val_.data());
+  sparse_arena_mat_t arena_adj_x(m, n, w.adj().size(), u_arena.data(),
+                                 v_arena.data(), w.vi_->adj_.data());
+  var_value<sparse_mat_t> var_x(arena_val_x, arena_adj_x);
+  return var_x;
+}
 
 /**
  * Create a sparse matrix from the given AoS matrix of vars and indexes.
  * @tparam Options Eigen matrix options.
- * @tparam MatrixVar A type inheriting from `Eigen::DenseBase` with a scalar type of @ref var_value
+ * @tparam MatrixVar A type inheriting from `Eigen::DenseBase` with a scalar
+ * type of @ref var_value
  * @tparam Vec1 Container type of the column indexes.
  * @tparam Vec2 Container type of the row indexes.
  * @param m Number of rows in matrix.
@@ -59,36 +59,41 @@ inline auto to_soa_sparse_matrix(int m, int n, VarMatrix&& w, Vec1&& u, Vec2&& v
  *          length as w.
  * @return Sparse matrix.
  */
-template <int Options = Eigen::ColMajor, typename MatrixVar, typename Vec1, typename Vec2,
-  require_eigen_dense_base_vt<is_var, MatrixVar>* = nullptr,
-  require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
-inline auto to_soa_sparse_matrix(int m, int n, MatrixVar&& w, Vec1&& u, Vec2&& v) {
-    auto w_arena = to_arena(std::forward<MatrixVar>(w));
-    auto u_arena = to_arena(std::forward<Vec1>(u));
-    auto v_arena = to_arena(std::forward<Vec2>(v));
-    arena_t<Eigen::SparseMatrix<var, Options>> arena_x(m, n, w_arena.size(),
-      u_arena.data(), v_arena.data(), w_arena.data());
-    var_value<Eigen::SparseMatrix<double, Options>> var_x(value_of(arena_x));
-    // No need to copy adj, but need to backprop
-    reverse_pass_callback([arena_x, var_x]() mutable {
-      using var_sparse_iterator_t = typename arena_t<Eigen::SparseMatrix<var, Options>>::InnerIterator;
-      using dbl_sparse_iterator_t = typename arena_t<Eigen::SparseMatrix<double, Options>>::InnerIterator;
-      // arena_x.adj() += var_x.adj() once custom adj() for var sparse matrix
-      for (int k = 0; k < arena_x.outerSize(); ++k) {
-        var_sparse_iterator_t it_arena_x(arena_x, k);
-        dbl_sparse_iterator_t it_var_x(var_x.adj(), k);
-        for (; bool(it_arena_x) && bool(it_var_x); ++it_arena_x, ++it_var_x) {
-          it_arena_x.valueRef().adj() += it_var_x.valueRef();
-        }
+template <int Options = Eigen::ColMajor, typename MatrixVar, typename Vec1,
+          typename Vec2,
+          require_eigen_dense_base_vt<is_var, MatrixVar>* = nullptr,
+          require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, MatrixVar&& w, Vec1&& u,
+                                 Vec2&& v) {
+  auto w_arena = to_arena(std::forward<MatrixVar>(w));
+  auto u_arena = to_arena(std::forward<Vec1>(u));
+  auto v_arena = to_arena(std::forward<Vec2>(v));
+  arena_t<Eigen::SparseMatrix<var, Options>> arena_x(
+      m, n, w_arena.size(), u_arena.data(), v_arena.data(), w_arena.data());
+  var_value<Eigen::SparseMatrix<double, Options>> var_x(value_of(arena_x));
+  // No need to copy adj, but need to backprop
+  reverse_pass_callback([arena_x, var_x]() mutable {
+    using var_sparse_iterator_t =
+        typename arena_t<Eigen::SparseMatrix<var, Options>>::InnerIterator;
+    using dbl_sparse_iterator_t =
+        typename arena_t<Eigen::SparseMatrix<double, Options>>::InnerIterator;
+    // arena_x.adj() += var_x.adj() once custom adj() for var sparse matrix
+    for (int k = 0; k < arena_x.outerSize(); ++k) {
+      var_sparse_iterator_t it_arena_x(arena_x, k);
+      dbl_sparse_iterator_t it_var_x(var_x.adj(), k);
+      for (; bool(it_arena_x) && bool(it_var_x); ++it_arena_x, ++it_var_x) {
+        it_arena_x.valueRef().adj() += it_var_x.valueRef();
       }
-    });
-    return var_x;
-  }
+    }
+  });
+  return var_x;
+}
 
 /**
  * Create a sparse matrix from the given matrix of floats and indexes.
  * @tparam Options Eigen matrix options.
- * @tparam Mat A type inheriting from `Eigen::DenseBase` with an arithmetic scalar type
+ * @tparam Mat A type inheriting from `Eigen::DenseBase` with an arithmetic
+ * scalar type
  * @tparam Vec1 Container type of the column indexes.
  * @tparam Vec2 Container type of the row indexes.
  * @param m Number of rows in matrix.
@@ -100,20 +105,20 @@ inline auto to_soa_sparse_matrix(int m, int n, MatrixVar&& w, Vec1&& u, Vec2&& v
  *          length as w.
  * @return Sparse matrix.
  */
-template <int Options = Eigen::ColMajor, typename Mat, typename Vec1, typename Vec2,
-  require_eigen_dense_base_vt<std::is_arithmetic, Mat>* = nullptr,
-  require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
-  inline auto to_soa_sparse_matrix(int m, int n, Mat&& w, Vec1&& u, Vec2&& v) {
-    auto w_arena = to_arena(std::forward<Mat>(w));
-    auto u_arena = to_arena(std::forward<Vec1>(u));
-    auto v_arena = to_arena(std::forward<Vec2>(v));
-    arena_t<Eigen::SparseMatrix<double, Options>> arena_x(m, n, w_arena.size(),
-      u_arena.data(), v_arena.data(), w_arena.data());
-    return var_value<Eigen::SparseMatrix<double, Options>>(arena_x);
-  }
-
-}
+template <int Options = Eigen::ColMajor, typename Mat, typename Vec1,
+          typename Vec2,
+          require_eigen_dense_base_vt<std::is_arithmetic, Mat>* = nullptr,
+          require_all_std_vector_vt<std::is_integral, Vec1, Vec2>* = nullptr>
+inline auto to_soa_sparse_matrix(int m, int n, Mat&& w, Vec1&& u, Vec2&& v) {
+  auto w_arena = to_arena(std::forward<Mat>(w));
+  auto u_arena = to_arena(std::forward<Vec1>(u));
+  auto v_arena = to_arena(std::forward<Vec2>(v));
+  arena_t<Eigen::SparseMatrix<double, Options>> arena_x(
+      m, n, w_arena.size(), u_arena.data(), v_arena.data(), w_arena.data());
+  return var_value<Eigen::SparseMatrix<double, Options>>(arena_x);
 }
 
-#endif
+}  // namespace math
+}  // namespace stan
 
+#endif
diff --git a/stan/math/rev/fun/to_var_value.hpp b/stan/math/rev/fun/to_var_value.hpp
index 6664b7d7521..4b93cf1e6e6 100644
--- a/stan/math/rev/fun/to_var_value.hpp
+++ b/stan/math/rev/fun/to_var_value.hpp
@@ -18,7 +18,8 @@ namespace math {
  * @param a matrix to convert
  */
 template <typename T, require_eigen_vt<is_var, T>* = nullptr>
-inline var_value<Eigen::Matrix<double, T::RowsAtCompileTime, T::ColsAtCompileTime>>
+inline var_value<
+    Eigen::Matrix<double, T::RowsAtCompileTime, T::ColsAtCompileTime>>
 to_var_value(const T& a) {
   arena_matrix<plain_type_t<T>> a_arena = a;
   var_value<promote_scalar_t<double, T>> res(a_arena.val());
diff --git a/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp b/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
index f17980e3d29..d17d1860fc3 100644
--- a/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
+++ b/test/unit/math/mix/fun/csr_matrix_times_vector_test.cpp
@@ -21,7 +21,6 @@ TEST(MathMixMatFun, csr_matrix_times_vector_vals) {
   EXPECT_MATRIX_EQ(dbl_res, var_res_bw);
 }
 
-
 TEST(MathMixMatFun, csr_matrix_times_vector1) {
   auto f = [](const auto& w, const auto& b) {
     using stan::math::csr_matrix_times_vector;
@@ -53,4 +52,3 @@ TEST(MathMixMatFun, csr_matrix_times_vector2) {
 
   stan::test::expect_ad_matvar(f, w, b);
 }
-
diff --git a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
index 70db6bc3571..a9feff41ce1 100644
--- a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
+++ b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
@@ -43,9 +43,9 @@ TEST_F(AgradRev, to_soa_sparse_matrix_matrix_var) {
 */
 
 TEST_F(AgradRev, to_soa_sparse_matrix_var_matrix) {
+  using stan::math::to_soa_sparse_matrix;
   using stan::math::var;
   using stan::math::var_value;
-  using stan::math::to_soa_sparse_matrix;
   std::vector<int> v{0, 1, 2, 0, 1};
   std::vector<int> u{0, 1, 2, 3, 4, 5};
   Eigen::VectorXd w(5);
@@ -53,8 +53,8 @@ TEST_F(AgradRev, to_soa_sparse_matrix_var_matrix) {
   int n = 5;
   w << 1, 2, 3, 4, 5;
   var_value<Eigen::VectorXd> w_var(w);
-  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena =
-    to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w_var, u, v);
+  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena
+      = to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w_var, u, v);
   for (int i = 0; i < 5; ++i) {
     EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w.val()(i));
   }
@@ -74,7 +74,4 @@ TEST_F(AgradRev, to_soa_sparse_matrix_var_matrix) {
   for (int i = 0; i < 5; ++i) {
     EXPECT_EQ(w_mat_arena.adj().valuePtr()[i], w_var.adj()(i));
   }
-
- }
-
-
+}

From f594b2cadc72124e05f6c5e59947e124f7861a2c Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 15 Apr 2024 16:24:37 -0400
Subject: [PATCH 5/9] fix header error

---
 stan/math/prim/meta/promote_scalar_type.hpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/stan/math/prim/meta/promote_scalar_type.hpp b/stan/math/prim/meta/promote_scalar_type.hpp
index da1a2deb405..fbcbbe1f749 100644
--- a/stan/math/prim/meta/promote_scalar_type.hpp
+++ b/stan/math/prim/meta/promote_scalar_type.hpp
@@ -4,6 +4,8 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/meta/is_eigen.hpp>
 #include <stan/math/prim/meta/is_var.hpp>
+#include <stan/math/prim/meta/is_eigen_dense_base.hpp>
+#include <stan/math/prim/meta/is_eigen_sparse_base.hpp>
 #include <vector>
 
 namespace stan {

From d1feb1955e05d3babb94ab6372a23f8c1c64cc72 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 15 Apr 2024 17:10:59 -0400
Subject: [PATCH 6/9] use static_cast for bool conversion in sparse matrix
 loops

---
 stan/math/rev/fun/to_soa_sparse_matrix.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stan/math/rev/fun/to_soa_sparse_matrix.hpp b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
index 2b87a4ea165..a31150cca84 100644
--- a/stan/math/rev/fun/to_soa_sparse_matrix.hpp
+++ b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
@@ -81,7 +81,7 @@ inline auto to_soa_sparse_matrix(int m, int n, MatrixVar&& w, Vec1&& u,
     for (int k = 0; k < arena_x.outerSize(); ++k) {
       var_sparse_iterator_t it_arena_x(arena_x, k);
       dbl_sparse_iterator_t it_var_x(var_x.adj(), k);
-      for (; bool(it_arena_x) && bool(it_var_x); ++it_arena_x, ++it_var_x) {
+      for (; static_cast<bool>(it_arena_x) && static_cast<bool>(it_var_x); ++it_arena_x, ++it_var_x) {
         it_arena_x.valueRef().adj() += it_var_x.valueRef();
       }
     }

From 33f0825af37e6141d81884101ded2da122e922e6 Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Mon, 15 Apr 2024 17:11:54 -0400
Subject: [PATCH 7/9] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 stan/math/rev/fun/to_soa_sparse_matrix.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/stan/math/rev/fun/to_soa_sparse_matrix.hpp b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
index a31150cca84..b5e872850a1 100644
--- a/stan/math/rev/fun/to_soa_sparse_matrix.hpp
+++ b/stan/math/rev/fun/to_soa_sparse_matrix.hpp
@@ -81,7 +81,8 @@ inline auto to_soa_sparse_matrix(int m, int n, MatrixVar&& w, Vec1&& u,
     for (int k = 0; k < arena_x.outerSize(); ++k) {
       var_sparse_iterator_t it_arena_x(arena_x, k);
       dbl_sparse_iterator_t it_var_x(var_x.adj(), k);
-      for (; static_cast<bool>(it_arena_x) && static_cast<bool>(it_var_x); ++it_arena_x, ++it_var_x) {
+      for (; static_cast<bool>(it_arena_x) && static_cast<bool>(it_var_x);
+           ++it_arena_x, ++it_var_x) {
         it_arena_x.valueRef().adj() += it_var_x.valueRef();
       }
     }

From 5c5dfc68b9641447bc5a325273c1e6090a0e0471 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Thu, 18 Apr 2024 12:25:44 -0400
Subject: [PATCH 8/9] update csr_matrix_times_vector w adjoint update.
 Uncomment tests for to_soa_sparse_matrix

---
 stan/math/rev/fun/csr_matrix_times_vector.hpp   | 17 ++---------------
 .../math/rev/fun/to_soa_sparse_matrix_test.cpp  |  4 ++--
 2 files changed, 4 insertions(+), 17 deletions(-)

diff --git a/stan/math/rev/fun/csr_matrix_times_vector.hpp b/stan/math/rev/fun/csr_matrix_times_vector.hpp
index 9af92b8d83a..4a0a93876e4 100644
--- a/stan/math/rev/fun/csr_matrix_times_vector.hpp
+++ b/stan/math/rev/fun/csr_matrix_times_vector.hpp
@@ -75,14 +75,7 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
         = 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 {
-      for (int k = 0; k < w_mat_arena.adj().outerSize(); ++k) {
-        for (typename sparse_var_value_t::vari_type::InnerIterator it(
-                 w_mat_arena.adj(), k);
-             it; ++it) {
-          it.valueRef()
-              += res.adj().coeff(it.row()) * b_arena.val().coeff(it.col());
-        }
-      }
+      w_mat_arena.adj() += res.adj() * b_arena.val().transpose();
       b_arena.adj() += w_mat_arena.val().transpose() * res.adj();
     });
     return return_t(res);
@@ -102,13 +95,7 @@ inline auto csr_matrix_times_vector(int m, int n, const T1& w,
     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 {
-      for (int k = 0; k < w_mat_arena.adj().outerSize(); ++k) {
-        for (typename sparse_var_value_t::vari_type::InnerIterator it(
-                 w_mat_arena.adj(), k);
-             it; ++it) {
-          it.valueRef() += res.adj().coeff(it.row()) * b_arena.coeff(it.col());
-        }
-      }
+      w_mat_arena.adj() += res.adj() * b_arena.transpose();
     });
     return return_t(res);
   }
diff --git a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
index a9feff41ce1..80ac4bd55cd 100644
--- a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
+++ b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
@@ -4,7 +4,7 @@
 #include <stan/math/rev/core.hpp>
 #include <test/unit/math/rev/util.hpp>
 #include <vector>
-/*
+
 TEST_F(AgradRev, to_soa_sparse_matrix_matrix_double) {
   using stan::math::var;
   using stan::math::var_value;
@@ -40,7 +40,7 @@ TEST_F(AgradRev, to_soa_sparse_matrix_matrix_var) {
     EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w.val()(i));
   }
 }
-*/
+
 
 TEST_F(AgradRev, to_soa_sparse_matrix_var_matrix) {
   using stan::math::to_soa_sparse_matrix;

From b0815c439e72a6ceec1eca5f6ccdd74620a7519e Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Thu, 18 Apr 2024 12:27:05 -0400
Subject: [PATCH 9/9] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 .../math/rev/fun/to_soa_sparse_matrix_test.cpp     | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
index 80ac4bd55cd..fbc495ac432 100644
--- a/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
+++ b/test/unit/math/rev/fun/to_soa_sparse_matrix_test.cpp
@@ -6,27 +6,26 @@
 #include <vector>
 
 TEST_F(AgradRev, to_soa_sparse_matrix_matrix_double) {
+  using stan::math::to_soa_sparse_matrix;
   using stan::math::var;
   using stan::math::var_value;
-  using stan::math::to_soa_sparse_matrix;
   std::vector<int> v{0, 1, 2, 0, 1};
   std::vector<int> u{0, 1, 2, 3, 4, 5};
   Eigen::VectorXd w(5);
   int m = 5;
   int n = 5;
   w << 1, 2, 3, 4, 5;
-  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena =
-    to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w, u, v);
+  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena
+      = to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w, u, v);
   for (int i = 0; i < 5; ++i) {
     EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w(i));
   }
 }
 
-
 TEST_F(AgradRev, to_soa_sparse_matrix_matrix_var) {
+  using stan::math::to_soa_sparse_matrix;
   using stan::math::var;
   using stan::math::var_value;
-  using stan::math::to_soa_sparse_matrix;
   std::vector<int> v{0, 1, 2, 0, 1};
   std::vector<int> u{0, 1, 2, 3, 4, 5};
   Eigen::VectorXd w(5);
@@ -34,14 +33,13 @@ TEST_F(AgradRev, to_soa_sparse_matrix_matrix_var) {
   int n = 5;
   w << 1, 2, 3, 4, 5;
   Eigen::Matrix<var, -1, 1> w_var(w);
-  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena =
-    to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w_var, u, v);
+  var_value<Eigen::SparseMatrix<double, Eigen::RowMajor>> w_mat_arena
+      = to_soa_sparse_matrix<Eigen::RowMajor>(m, n, w_var, u, v);
   for (int i = 0; i < 5; ++i) {
     EXPECT_EQ(w_mat_arena.val().valuePtr()[i], w.val()(i));
   }
 }
 
-
 TEST_F(AgradRev, to_soa_sparse_matrix_var_matrix) {
   using stan::math::to_soa_sparse_matrix;
   using stan::math::var;