From 84ce960df000a67892e382ce71af5fc30e9cdf68 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Tue, 7 Nov 2023 17:54:12 -0500
Subject: [PATCH 01/14] Updates vari_value<SparseMatrix> to use a new
 arena_matrix<SparseMatrix> as it's inner matrix type

---
 stan/math/rev/core/arena_matrix.hpp           | 132 +++++++++++++++++-
 stan/math/rev/core/vari.hpp                   |  12 +-
 stan/math/rev/meta/arena_type.hpp             |   2 +-
 test/unit/math/rev/core/arena_matrix_test.cpp |  84 ++++++++++-
 test/unit/util.hpp                            |   1 +
 5 files changed, 217 insertions(+), 14 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 3d3b095fc74..2fb8b6d9186 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -4,7 +4,7 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
 #include <stan/math/rev/core/chainablestack.hpp>
-
+#include <iostream>
 namespace stan {
 namespace math {
 
@@ -15,8 +15,11 @@ namespace math {
  * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`
  * ...)
  */
+template <typename MatrixType, typename>
+class arena_matrix;
+
 template <typename MatrixType>
-class arena_matrix : public Eigen::Map<MatrixType> {
+class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>> : public Eigen::Map<MatrixType> {
  public:
   using Scalar = value_type_t<MatrixType>;
   using Base = Eigen::Map<MatrixType>;
@@ -130,6 +133,131 @@ class arena_matrix : public Eigen::Map<MatrixType> {
   }
 };
 
+template <typename MatrixType>
+class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>> : public Eigen::Map<MatrixType> {
+ public:
+  using Scalar = value_type_t<MatrixType>;
+  using Base = Eigen::Map<MatrixType>;
+  using PlainObject = std::decay_t<MatrixType>;
+  using StorageIndex = typename PlainObject::StorageIndex;
+  using storage_idx_t = arena_matrix<Eigen::Matrix<StorageIndex, Eigen::Dynamic, 1>>;
+  using data_t = arena_matrix<Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>;
+  /**
+   * Default constructor.
+   */
+  arena_matrix()
+      : Base::Map(0, 0, 0, nullptr, nullptr, nullptr) {
+  }
+
+private:
+  template <typename T, typename Integral>
+  auto* copy_vector(T* ptr, Integral size) {
+    arena_matrix<Eigen::Matrix<T, Eigen::Dynamic, 1>> ret(size);
+    std::copy_n(ptr, size, ret.data());
+    return ret.data();
+  }
+public:
+  arena_matrix(Eigen::Index rows, Eigen::Index cols, Eigen::Index nnz,
+   StorageIndex* outerIndexPtr, StorageIndex* innerIndexPtr,
+   Scalar* valuePtr, StorageIndex* innerNonZerosPtr = nullptr) :
+   Base::Map(rows, cols, nnz, outerIndexPtr, innerIndexPtr, valuePtr, innerNonZerosPtr) {}
+  /**
+   * Constructs `arena_matrix` from an expression.
+   * @param other expression
+   */
+  template <typename T, require_same_t<T, PlainObject>* = nullptr>
+  arena_matrix(T&& other)  // NOLINT
+      :  Base::Map(other.rows(), other.cols(), other.nonZeros(),
+       copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
+       copy_vector(other.innerIndexPtr(), other.nonZeros()),
+       copy_vector(other.valuePtr(), other.nonZeros())) {
+  }
+
+  template <typename S, require_convertible_t<S&, PlainObject>* = nullptr,
+   require_not_same_t<S, PlainObject>* = nullptr>
+  arena_matrix(S&& x) : arena_matrix(PlainObject(x)) {}
+
+  /**
+   * Constructs `arena_matrix` from an expression. This makes an assumption that
+   * any other `Eigen::Map` also contains memory allocated in the arena.
+   * @param other expression
+   */
+  arena_matrix(const Base& other)  // NOLINT
+      : Base(other) {}
+
+  /**
+   * Copy constructor.
+   * @param other matrix to copy from
+   */
+  arena_matrix(const arena_matrix<MatrixType>& other)
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+         other.outerIndexPtr(), other.innerIndexPtr(), other.valuePtr()) {}
+
+
+  // without this using, compiler prefers combination of implicit construction
+  // and copy assignment to the inherited operator when assigned an expression
+  using Base::operator=;
+
+  /**
+   * Copy assignment operator.
+   * @param other matrix to copy from
+   * @return `*this`
+   */
+  arena_matrix& operator=(const arena_matrix<MatrixType>& other) {
+    // placement new changes what data map points to - there is no allocation
+    new (this)
+        Base(other.rows(), other.cols(), other.nonZeros(),
+        other.outerIndexPtr(), other.innerIndexPtr(), other.valuePtr());
+    return *this;
+  }
+      /**
+   * Assignment operator for assigning an expression.
+   * @param a expression to evaluate into this
+   * @return `*this`
+   */
+  template <typename T>
+  arena_matrix& operator=(const T& expr) {
+    PlainObject x = expr;
+    const auto alloc_storage = [](auto size) {
+      ChainableStack::instance_->memalloc_.alloc_array<StorageIndex>(size);
+    };
+    std::cout << "Got1" << std::endl;
+    new (this) Base(x.rows(), x.cols(), x.nonZeros(), alloc_storage(x.innerSize()),
+      alloc_storage(x.outerSize()),
+      ChainableStack::instance_->memalloc_.alloc_array<Scalar>(x.nonZeros()));
+    std::cout << "Got2" << std::endl;
+    Base::operator=(x);
+    std::cout << "Got3" << std::endl;
+    return *this;
+  }
+
+  arena_matrix& operator+=(const arena_matrix<MatrixType>& other) {
+    using inner_iterator = typename Base::InnerIterator;
+    for (int k = 0; k < (*this).outerSize(); ++k) {
+      for (inner_iterator it(*this, k), iz(other, k); it; ++it, ++iz) {
+        iz.value() += it.value();
+      }
+    }
+    return *this;
+  }
+
+  template <typename Expr, require_convertible_t<Expr&, MatrixType>* = nullptr,
+   require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  arena_matrix& operator+=(Expr&& other) {
+    using inner_iterator = typename Base::InnerIterator;
+    auto&& x = to_ref(other);
+    for (int k = 0; k < (*this).outerSize(); ++k) {
+      inner_iterator it(*this, k);
+      typename Expr::InnerIterator iz(x, k);
+      for (; it; ++it, ++iz) {
+        iz.value() += it.value();
+      }
+    }
+    return *this;
+  }
+
+};
+
 }  // namespace math
 }  // namespace stan
 
diff --git a/stan/math/rev/core/vari.hpp b/stan/math/rev/core/vari.hpp
index f61d2917dfc..0ab09942d5e 100644
--- a/stan/math/rev/core/vari.hpp
+++ b/stan/math/rev/core/vari.hpp
@@ -807,8 +807,7 @@ class vari_value<T, require_all_t<is_plain_type<T>, is_eigen_dense_base<T>>>
  *
  */
 template <typename T>
-class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
-                                                      chainable_alloc {
+class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
  public:
   using PlainObject = plain_type_t<T>;  // Base type of Eigen class
   using value_type = PlainObject;       // vari's adj_ and val_ member type
@@ -825,12 +824,12 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    * The adjoint of this variable, which is the partial derivative
    * of this variable with respect to the root variable.
    */
-  PlainObject adj_;
+  arena_matrix<PlainObject> adj_;
 
   /**
    * The value of this variable.
    */
-  const PlainObject val_;
+  arena_matrix<PlainObject> val_;
 
   /**
    * Construct a variable implementation from a value. The
@@ -847,8 +846,7 @@ 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)
-      : chainable_alloc(), adj_(x), val_(std::forward<S>(x)) {
+  explicit vari_value(S&& x) : adj_(x), val_(std::forward<S>(x)) {
     this->set_zero_adjoint();
     ChainableStack::instance_->var_stack_.push_back(this);
   }
@@ -871,7 +869,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base,
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
   vari_value(S&& x, bool stacked)
-      : chainable_alloc(), adj_(x), val_(std::forward<S>(x)) {
+      : adj_(x), val_(std::forward<S>(x)) {
     this->set_zero_adjoint();
     if (stacked) {
       ChainableStack::instance_->var_stack_.push_back(this);
diff --git a/stan/math/rev/meta/arena_type.hpp b/stan/math/rev/meta/arena_type.hpp
index db175293bc5..28c7248b6e5 100644
--- a/stan/math/rev/meta/arena_type.hpp
+++ b/stan/math/rev/meta/arena_type.hpp
@@ -10,7 +10,7 @@
 namespace stan {
 namespace math {
 // forward declaration
-template <typename MatrixType>
+template <typename MatrixType, typename = void>
 class arena_matrix;
 }  // namespace math
 
diff --git a/test/unit/math/rev/core/arena_matrix_test.cpp b/test/unit/math/rev/core/arena_matrix_test.cpp
index 1bc59cac007..ba97909453e 100644
--- a/test/unit/math/rev/core/arena_matrix_test.cpp
+++ b/test/unit/math/rev/core/arena_matrix_test.cpp
@@ -2,7 +2,7 @@
 #include <test/unit/util.hpp>
 #include <gtest/gtest.h>
 
-TEST(AgradRevArenaMat, arena_matrix_matrix_test) {
+TEST(AgradRev, arena_matrix_matrix_test) {
   using Eigen::MatrixXd;
   using stan::math::arena_matrix;
 
@@ -29,7 +29,7 @@ TEST(AgradRevArenaMat, arena_matrix_matrix_test) {
   stan::math::recover_memory();
 }
 
-TEST(AgradRevArenaMat, arena_matrix_vector_test) {
+TEST(AgradRev, arena_matrix_vector_test) {
   using Eigen::VectorXd;
   using stan::math::arena_matrix;
 
@@ -56,7 +56,7 @@ TEST(AgradRevArenaMat, arena_matrix_vector_test) {
   stan::math::recover_memory();
 }
 
-TEST(AgradRevArenaMat, arena_matrix_row_vector_test) {
+TEST(AgradRev, arena_matrix_row_vector_test) {
   using Eigen::RowVectorXd;
   using stan::math::arena_matrix;
 
@@ -83,7 +83,7 @@ TEST(AgradRevArenaMat, arena_matrix_row_vector_test) {
   stan::math::recover_memory();
 }
 
-TEST(AgradRevArenaMat, arena_matrix_transpose_test) {
+TEST(AgradRev, arena_matrix_transpose_test) {
   using stan::math::arena_matrix;
 
   Eigen::VectorXd c = Eigen::VectorXd::Random(3);
@@ -104,3 +104,79 @@ TEST(AgradRevArenaMat, arena_matrix_transpose_test) {
 
   stan::math::recover_memory();
 }
+
+template <typename T_map, typename T_mat>
+void expect_sparse_matrix_equal(T_map&& map_x, T_mat&& mat_x) {
+  using inner_iterator = typename std::decay_t<T_mat>::InnerIterator;
+  using map_inner_iterator = typename std::decay_t<T_map>::InnerIterator;
+  Eigen::Index nnz_m = map_x.nonZeros();
+  Eigen::Index nnz = mat_x.nonZeros();
+  EXPECT_EQ(map_x.nonZeros(), mat_x.nonZeros());
+  for (int k = 0; k < mat_x.outerSize(); ++k) {
+    inner_iterator it(mat_x, k);
+    map_inner_iterator iz(map_x, k);
+    for (; it && iz; ++it, ++iz) {
+      EXPECT_FLOAT_EQ(iz.value(), it.value());
+    }
+  }
+}
+
+TEST(AgradRev, arena_sparse_matrix_constructors) {
+  using eig_mat = Eigen::SparseMatrix<double>;
+  using stan::math::arena_matrix;
+  using arena_mat = arena_matrix<eig_mat>;
+  using inner_iterator = typename eig_mat::InnerIterator;
+  using map_inner_iterator = typename Eigen::Map<eig_mat>::InnerIterator;
+  using stan::test::make_sparse_matrix_random;
+  eig_mat A = make_sparse_matrix_random(10, 10);
+  // Testing each constructor
+  arena_mat empty_m();
+  arena_mat nocopy(A.rows(), A.cols(), A.nonZeros(),
+    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
+    A.innerNonZeroPtr());
+  Eigen::Map<eig_mat> sparse_map(A.rows(), A.cols(), A.nonZeros(),
+    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
+    A.innerNonZeroPtr());
+  arena_mat sparse_map_constructor(sparse_map);
+  arena_mat A_m(A);
+  expect_sparse_matrix_equal(A_m, A);
+  arena_mat arena_mat_constructor(A_m);
+  expect_sparse_matrix_equal(arena_mat_constructor, A);
+  eig_mat mul_A = A * A;
+  arena_mat mul_A_m(A * A_m);
+  expect_sparse_matrix_equal(mul_A_m, mul_A);
+  eig_mat mul_B = mul_A * A;
+  arena_mat mul_B_m(mul_A_m * A_m);
+  expect_sparse_matrix_equal(mul_B_m, mul_B);
+
+}
+
+TEST(AgradRev, arena_sparse_matrix_equals) {
+  using eig_mat = Eigen::SparseMatrix<double>;
+  using stan::math::arena_matrix;
+  using arena_mat = arena_matrix<eig_mat>;
+  using inner_iterator = typename eig_mat::InnerIterator;
+  using map_inner_iterator = typename Eigen::Map<eig_mat>::InnerIterator;
+  using stan::test::make_sparse_matrix_random;
+  eig_mat A = make_sparse_matrix_random(10, 10);
+  // Testing each constructor
+  arena_mat empty_m = arena_mat{};
+  arena_mat nocopy = arena_mat(A.rows(), A.cols(), A.nonZeros(),
+    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
+    A.innerNonZeroPtr());
+  Eigen::Map<eig_mat> sparse_map(A.rows(), A.cols(), A.nonZeros(),
+    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
+    A.innerNonZeroPtr());
+  arena_mat sparse_map_constructor = sparse_map;
+  arena_mat A_m = A;
+  expect_sparse_matrix_equal(A_m, A);
+  arena_mat arena_mat_constructor = A_m;
+  expect_sparse_matrix_equal(arena_mat_constructor, A);
+  eig_mat mul_A = A * A;
+  arena_mat mul_A_m = A * A_m;
+  expect_sparse_matrix_equal(mul_A_m, mul_A);
+  eig_mat mul_B = mul_A * A;
+  arena_mat mul_B_m = mul_A_m * A_m;
+  expect_sparse_matrix_equal(mul_B_m, mul_B);
+
+}
diff --git a/test/unit/util.hpp b/test/unit/util.hpp
index 554d4d90556..6d13989d579 100644
--- a/test/unit/util.hpp
+++ b/test/unit/util.hpp
@@ -227,6 +227,7 @@ auto make_sparse_matrix_random(int rows, int cols) {
   }
   Eigen::SparseMatrix<double> mat(rows, cols);
   mat.setFromTriplets(tripletList.begin(), tripletList.end());
+  mat.makeCompressed();
   return mat;
 }
 

From 9890d1a9cc82842ad2d81a6b5ab08acfeb482197 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Wed, 8 Nov 2023 15:21:31 -0500
Subject: [PATCH 02/14] Adds arena_matrix specialization for sparse matrices

---
 stan/math/rev/core/arena_matrix.hpp           | 257 +++++++++++++-----
 stan/math/rev/core/vari.hpp                   |  13 +-
 test/unit/math/rev/core/arena_matrix_test.cpp | 102 +++++--
 test/unit/math/rev/core/var_test.cpp          |  10 +-
 test/unit/math/rev/core/vari_test.cpp         |   3 +-
 test/unit/math/rev/util.hpp                   |   7 +
 6 files changed, 287 insertions(+), 105 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 2fb8b6d9186..61aca79b3c1 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -4,6 +4,7 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
 #include <stan/math/rev/core/chainablestack.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
 #include <iostream>
 namespace stan {
 namespace math {
@@ -19,7 +20,8 @@ template <typename MatrixType, typename>
 class arena_matrix;
 
 template <typename MatrixType>
-class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>> : public Eigen::Map<MatrixType> {
+class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>>
+    : public Eigen::Map<MatrixType> {
  public:
   using Scalar = value_type_t<MatrixType>;
   using Base = Eigen::Map<MatrixType>;
@@ -134,48 +136,72 @@ class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>> : public
 };
 
 template <typename MatrixType>
-class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>> : public Eigen::Map<MatrixType> {
+class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
+    : public Eigen::Map<MatrixType> {
  public:
   using Scalar = value_type_t<MatrixType>;
   using Base = Eigen::Map<MatrixType>;
   using PlainObject = std::decay_t<MatrixType>;
   using StorageIndex = typename PlainObject::StorageIndex;
-  using storage_idx_t = arena_matrix<Eigen::Matrix<StorageIndex, Eigen::Dynamic, 1>>;
-  using data_t = arena_matrix<Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>;
+
   /**
    * Default constructor.
    */
-  arena_matrix()
-      : Base::Map(0, 0, 0, nullptr, nullptr, nullptr) {
-  }
+  arena_matrix() : Base::Map(0, 0, 0, nullptr, nullptr, nullptr) {}
+  /**
+   * Constructor for CDC formatted data. Assumes compressed and does not own
+   * memory.
+   *
+   * Constructs a read-write Map to a sparse matrix of size rows x cols,
+   * containing nnz non-zero coefficients, stored as a sparse format as defined
+   * by the pointers outerIndexPtr, innerIndexPtr, and valuePtr. If the optional
+   * parameter innerNonZerosPtr is the null pointer, then a standard compressed
+   * format is assumed.
+   *
+   * @param rows Number of rows
+   * @param cols Number of columns
+   * @param nnz Number of nonzero items in matrix
+   * @param outerIndexPtr A pointer to the array of outer indices.
+   * @param innerIndexPtr A pointer to the array of inner indices.
+   * @param valuePtr A pointer to the array of values.
+   * @param innerNonZerosPtr A pointer to the array of the number of non zeros
+   * of the inner vectors.
+   *
+   */
+  arena_matrix(Eigen::Index rows, Eigen::Index cols, Eigen::Index nnz,
+               StorageIndex* outerIndexPtr, StorageIndex* innerIndexPtr,
+               Scalar* valuePtr, StorageIndex* innerNonZerosPtr = nullptr)
+      : Base::Map(rows, cols, nnz, outerIndexPtr, innerIndexPtr, valuePtr,
+                  innerNonZerosPtr) {}
 
-private:
+ private:
   template <typename T, typename Integral>
-  auto* copy_vector(T* ptr, Integral size) {
-    arena_matrix<Eigen::Matrix<T, Eigen::Dynamic, 1>> ret(size);
-    std::copy_n(ptr, size, ret.data());
-    return ret.data();
+  auto* copy_vector(const T* ptr, Integral size) {
+    T* ret = ChainableStack::instance_->memalloc_.alloc_array<T>(size);
+    std::copy_n(ptr, size, ret);
+    return ret;
   }
-public:
-  arena_matrix(Eigen::Index rows, Eigen::Index cols, Eigen::Index nnz,
-   StorageIndex* outerIndexPtr, StorageIndex* innerIndexPtr,
-   Scalar* valuePtr, StorageIndex* innerNonZerosPtr = nullptr) :
-   Base::Map(rows, cols, nnz, outerIndexPtr, innerIndexPtr, valuePtr, innerNonZerosPtr) {}
+
+ public:
   /**
-   * Constructs `arena_matrix` from an expression.
-   * @param other expression
+   * Constructs `arena_matrix` from an `Eigen::SparseMatrix`.
+   * @param other Eigen Sparse Matrix class
    */
   template <typename T, require_same_t<T, PlainObject>* = nullptr>
   arena_matrix(T&& other)  // NOLINT
-      :  Base::Map(other.rows(), other.cols(), other.nonZeros(),
-       copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
-       copy_vector(other.innerIndexPtr(), other.nonZeros()),
-       copy_vector(other.valuePtr(), other.nonZeros())) {
-  }
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
+                  copy_vector(other.innerIndexPtr(), other.nonZeros()),
+                  copy_vector(other.valuePtr(), other.nonZeros())) {}
 
+  /**
+   * Constructs `arena_matrix` from an Eigen expression
+   * @param other An expression
+   */
   template <typename S, require_convertible_t<S&, PlainObject>* = nullptr,
-   require_not_same_t<S, PlainObject>* = nullptr>
-  arena_matrix(S&& x) : arena_matrix(PlainObject(x)) {}
+            require_not_same_t<S, PlainObject>* = nullptr>
+  arena_matrix(S&& x)  // NOLINT
+      : arena_matrix(PlainObject(x)) {}
 
   /**
    * Constructs `arena_matrix` from an expression. This makes an assumption that
@@ -186,13 +212,35 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>> : public
       : Base(other) {}
 
   /**
-   * Copy constructor.
+   * Const copy constructor. No actual copy is performed
+   * @note Since the memory for the arena matrix sits in Stan's memory arena all
+   * copies/moves of arena matrices are shallow moves
    * @param other matrix to copy from
    */
   arena_matrix(const arena_matrix<MatrixType>& other)
       : Base::Map(other.rows(), other.cols(), other.nonZeros(),
-         other.outerIndexPtr(), other.innerIndexPtr(), other.valuePtr()) {}
-
+                  other.outerIndexPtr(), other.innerIndexPtr(),
+                  other.valuePtr()) {}
+  /**
+   * Move constructor.
+   * @note Since the memory for the arena matrix sits in Stan's memory arena all
+   * copies/moves of arena matrices are shallow moves
+   * @param other matrix to copy from
+   */
+  arena_matrix(arena_matrix<MatrixType>&& other)
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  other.outerIndexPtr(), other.innerIndexPtr(),
+                  other.valuePtr()) {}
+  /**
+   * Copy constructor. No actual copy is performed
+   * @note Since the memory for the arena matrix sits in Stan's memory arena all
+   * copies/moves of arena matrices are shallow moves
+   * @param other matrix to copy from
+   */
+  arena_matrix(arena_matrix<MatrixType>& other)
+      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
+                  other.outerIndexPtr(), other.innerIndexPtr(),
+                  other.valuePtr()) {}
 
   // without this using, compiler prefers combination of implicit construction
   // and copy assignment to the inherited operator when assigned an expression
@@ -200,62 +248,149 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>> : public
 
   /**
    * Copy assignment operator.
+   * @tparam ArenaMatrix An `arena_matrix` type
    * @param other matrix to copy from
    * @return `*this`
    */
-  arena_matrix& operator=(const arena_matrix<MatrixType>& other) {
+  template <typename ArenaMatrix,
+            require_same_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(),
-        other.outerIndexPtr(), other.innerIndexPtr(), other.valuePtr());
+    new (this) Base(other.rows(), other.cols(), other.nonZeros(),
+                    const_cast<StorageIndex*>(other.outerIndexPtr()),
+                    const_cast<StorageIndex*>(other.innerIndexPtr()),
+                    const_cast<Scalar*>(other.valuePtr()));
     return *this;
   }
-      /**
+
+  /**
    * Assignment operator for assigning an expression.
-   * @param a expression to evaluate into this
+   * @tparam Expr An expression that an `arena_matrix<MatrixType>` can be
+   * constructed from
+   * @param expr expression to evaluate into this
    * @return `*this`
    */
-  template <typename T>
-  arena_matrix& operator=(const T& expr) {
-    PlainObject x = expr;
-    const auto alloc_storage = [](auto size) {
-      ChainableStack::instance_->memalloc_.alloc_array<StorageIndex>(size);
-    };
-    std::cout << "Got1" << std::endl;
-    new (this) Base(x.rows(), x.cols(), x.nonZeros(), alloc_storage(x.innerSize()),
-      alloc_storage(x.outerSize()),
-      ChainableStack::instance_->memalloc_.alloc_array<Scalar>(x.nonZeros()));
-    std::cout << "Got2" << std::endl;
-    Base::operator=(x);
-    std::cout << "Got3" << std::endl;
+  template <typename Expr,
+            require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  arena_matrix& operator=(Expr&& expr) {
+    *this = arena_matrix(std::forward<Expr>(expr));
     return *this;
   }
 
-  arena_matrix& operator+=(const arena_matrix<MatrixType>& other) {
-    using inner_iterator = typename Base::InnerIterator;
+ private:
+  /**
+   * inplace operations functor for `Sparse (.*)= Sparse`.
+   * @note This assumes that each matrix is of the same size and sparsity
+   * pattern.
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam Expr Type derived from `Eigen::SparseMatrixBase`
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename Expr,
+            require_convertible_t<Expr&, MatrixType>* = nullptr>
+  void inplace_ops_impl(F&& f, Expr&& other) {
+    auto&& x = to_ref(other);
     for (int k = 0; k < (*this).outerSize(); ++k) {
-      for (inner_iterator it(*this, k), iz(other, k); it; ++it, ++iz) {
-        iz.value() += it.value();
+      typename Base::InnerIterator it(*this, k);
+      typename std::decay_t<Expr>::InnerIterator iz(x, k);
+      for (; bool(it) && bool(iz); ++it, ++iz) {
+        f(it.valueRef(), iz.value());
       }
     }
-    return *this;
   }
 
-  template <typename Expr, require_convertible_t<Expr&, MatrixType>* = nullptr,
-   require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
-  arena_matrix& operator+=(Expr&& other) {
-    using inner_iterator = typename Base::InnerIterator;
+  /**
+   * inplace operations functor for `Sparse (.*)= Dense`.
+   * @note This assumes the user intends to perform the inplace operation for
+   * the nonzero parts of `this`
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam Expr Type derived from `Eigen::DenseBase`
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename Expr,
+            require_not_convertible_t<Expr&, MatrixType>* = nullptr,
+            require_eigen_dense_base_t<Expr>* = nullptr>
+  void inplace_ops_impl(F&& f, Expr&& other) {
     auto&& x = to_ref(other);
     for (int k = 0; k < (*this).outerSize(); ++k) {
-      inner_iterator it(*this, k);
-      typename Expr::InnerIterator iz(x, k);
-      for (; it; ++it, ++iz) {
-        iz.value() += it.value();
+      typename Base::InnerIterator it(*this, k);
+      for (; bool(it); ++it) {
+        f(it.valueRef(), x(it.row(), it.col()));
       }
     }
+  }
+
+  /**
+   * inplace operations functor for `Sparse (.*)= Scalar`.
+   * @note This assumes the user intends to perform the inplace operation for
+   * the nonzero parts of `this`
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam T A scalar type
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename T,
+            require_convertible_t<T&, Scalar>* = nullptr>
+  void inplace_ops_impl(F&& f, T&& other) {
+    for (int k = 0; k < (*this).outerSize(); ++k) {
+      typename Base::InnerIterator it(*this, k);
+      for (; bool(it); ++it) {
+        f(it.valueRef(), other);
+      }
+    }
+  }
+
+ public:
+  /**
+   * Inplace operator `+=`
+   * @note Caution!! Inplace operators assume that either
+   *  1. The right hand side sparse matrix has the same sparcity pattern
+   *  2. You only intend to add a scalar or dense matrix coefficients to the
+   * nonzero values of `this` This is intended to be used within the reverse
+   * pass for accumulation of the adjoint and is built as such. Any other use
+   * case should be be sure the above assumptions are satisfied.
+   * @tparam T A type derived from `Eigen::SparseMatrixBase` or
+   * `Eigen::DenseMatrixBase` or a `Scalar`
+   * @param other value to be added inplace to the matrix.
+   */
+  template <typename T>
+  arena_matrix& operator+=(T&& other) {
+    inplace_ops_impl(
+        [](auto&& x, auto&& y) {
+          x += y;
+          return;
+        },
+        std::forward<T>(other));
     return *this;
   }
 
+  /**
+   * Inplace operator `+=`
+   * @note Caution!! Inplace operators assume that either
+   *  1. The right hand side sparse matrix has the same sparcity pattern
+   *  2. You only intend to add a scalar or dense matrix coefficients to the
+   * nonzero values of `this` This is intended to be used within the reverse
+   * pass for accumulation of the adjoint and is built as such. Any other use
+   * case should be be sure the above assumptions are satisfied.
+   * @tparam T A type derived from `Eigen::SparseMatrixBase` or
+   * `Eigen::DenseMatrixBase` or a `Scalar`
+   * @param other value to be added inplace to the matrix.
+   */
+  template <typename T>
+  arena_matrix& operator-=(T&& other) {
+    inplace_ops_impl(
+        [](auto&& x, auto&& y) {
+          x -= y;
+          return;
+        },
+        std::forward<T>(other));
+    return *this;
+  }
 };
 
 }  // namespace math
diff --git a/stan/math/rev/core/vari.hpp b/stan/math/rev/core/vari.hpp
index 0ab09942d5e..f3caf6f925a 100644
--- a/stan/math/rev/core/vari.hpp
+++ b/stan/math/rev/core/vari.hpp
@@ -811,6 +811,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
  public:
   using PlainObject = plain_type_t<T>;  // Base type of Eigen class
   using value_type = PlainObject;       // vari's adj_ and val_ member type
+  using InnerIterator = typename arena_matrix<PlainObject>::InnerIterator;
   /**
    * Rows at compile time
    */
@@ -919,11 +920,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    * result with respect to itself to be 1.
    */
   inline void init_dependent() {
-    for (int k = 0; k < adj_.outerSize(); ++k) {
-      for (typename PlainObject::InnerIterator it(adj_, k); it; ++it) {
-        it.valueRef() = 1.0;
-      }
-    }
+    std::fill(adj_.valuePtr(), adj_.valuePtr() + adj_.nonZeros(), 1.0);
   }
 
   /**
@@ -932,11 +929,7 @@ class vari_value<T, require_eigen_sparse_base_t<T>> : public vari_base {
    * example in a Jacobian calculation).
    */
   inline void set_zero_adjoint() noexcept final {
-    for (int k = 0; k < adj_.outerSize(); ++k) {
-      for (typename PlainObject::InnerIterator it(adj_, k); it; ++it) {
-        it.valueRef() = 0.0;
-      }
-    }
+    std::fill(adj_.valuePtr(), adj_.valuePtr() + adj_.nonZeros(), 0.0);
   }
 
   /**
diff --git a/test/unit/math/rev/core/arena_matrix_test.cpp b/test/unit/math/rev/core/arena_matrix_test.cpp
index ba97909453e..5ed3ac0a6c2 100644
--- a/test/unit/math/rev/core/arena_matrix_test.cpp
+++ b/test/unit/math/rev/core/arena_matrix_test.cpp
@@ -1,8 +1,9 @@
 #include <stan/math.hpp>
 #include <test/unit/util.hpp>
+#include <test/unit/math/rev/util.hpp>
 #include <gtest/gtest.h>
 
-TEST(AgradRev, arena_matrix_matrix_test) {
+TEST_F(AgradRev, arena_matrix_matrix_test) {
   using Eigen::MatrixXd;
   using stan::math::arena_matrix;
 
@@ -26,10 +27,9 @@ TEST(AgradRev, arena_matrix_matrix_test) {
   a = MatrixXd::Ones(3, 2);
 
   EXPECT_MATRIX_EQ(a + a2 + a3 + b + b2 + e, MatrixXd::Ones(3, 2) * 9);
-  stan::math::recover_memory();
 }
 
-TEST(AgradRev, arena_matrix_vector_test) {
+TEST_F(AgradRev, arena_matrix_vector_test) {
   using Eigen::VectorXd;
   using stan::math::arena_matrix;
 
@@ -53,10 +53,9 @@ TEST(AgradRev, arena_matrix_vector_test) {
   a = VectorXd::Ones(3);
 
   EXPECT_MATRIX_EQ(a + a2 + a3 + b + b2 + e, VectorXd::Ones(3) * 9);
-  stan::math::recover_memory();
 }
 
-TEST(AgradRev, arena_matrix_row_vector_test) {
+TEST_F(AgradRev, arena_matrix_row_vector_test) {
   using Eigen::RowVectorXd;
   using stan::math::arena_matrix;
 
@@ -80,10 +79,9 @@ TEST(AgradRev, arena_matrix_row_vector_test) {
   a = RowVectorXd::Ones(3);
 
   EXPECT_MATRIX_EQ(a + a2 + a3 + b + b2 + e, RowVectorXd::Ones(3) * 9);
-  stan::math::recover_memory();
 }
 
-TEST(AgradRev, arena_matrix_transpose_test) {
+TEST_F(AgradRev, arena_matrix_transpose_test) {
   using stan::math::arena_matrix;
 
   Eigen::VectorXd c = Eigen::VectorXd::Random(3);
@@ -101,18 +99,14 @@ TEST(AgradRev, arena_matrix_transpose_test) {
   a = Eigen::VectorXd::Random(3);
   EXPECT_NO_THROW(b = a);
   EXPECT_MATRIX_EQ(a, b.transpose());
-
-  stan::math::recover_memory();
 }
 
 template <typename T_map, typename T_mat>
-void expect_sparse_matrix_equal(T_map&& map_x, T_mat&& mat_x) {
+inline void expect_sparse_matrix_equal(T_map&& map_x, T_mat&& mat_x) {
   using inner_iterator = typename std::decay_t<T_mat>::InnerIterator;
   using map_inner_iterator = typename std::decay_t<T_map>::InnerIterator;
-  Eigen::Index nnz_m = map_x.nonZeros();
-  Eigen::Index nnz = mat_x.nonZeros();
   EXPECT_EQ(map_x.nonZeros(), mat_x.nonZeros());
-  for (int k = 0; k < mat_x.outerSize(); ++k) {
+  for (int k = 0; k < map_x.outerSize(); ++k) {
     inner_iterator it(mat_x, k);
     map_inner_iterator iz(map_x, k);
     for (; it && iz; ++it, ++iz) {
@@ -121,7 +115,19 @@ void expect_sparse_matrix_equal(T_map&& map_x, T_mat&& mat_x) {
   }
 }
 
-TEST(AgradRev, arena_sparse_matrix_constructors) {
+template <typename T_map>
+inline void expect_sparse_dense_matrix_equal(
+    T_map&& map_x, const Eigen::Matrix<double, -1, -1>& mat_x) {
+  using map_inner_iterator = typename std::decay_t<T_map>::InnerIterator;
+  Eigen::Index nnz_m = map_x.nonZeros();
+  for (int k = 0; k < map_x.outerSize(); ++k) {
+    for (map_inner_iterator iz(map_x, k); iz; ++iz) {
+      EXPECT_FLOAT_EQ(iz.value(), mat_x(iz.row(), iz.col()));
+    }
+  }
+}
+
+TEST_F(AgradRev, arena_sparse_matrix_constructors) {
   using eig_mat = Eigen::SparseMatrix<double>;
   using stan::math::arena_matrix;
   using arena_mat = arena_matrix<eig_mat>;
@@ -131,12 +137,11 @@ TEST(AgradRev, arena_sparse_matrix_constructors) {
   eig_mat A = make_sparse_matrix_random(10, 10);
   // Testing each constructor
   arena_mat empty_m();
-  arena_mat nocopy(A.rows(), A.cols(), A.nonZeros(),
-    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
-    A.innerNonZeroPtr());
+  arena_mat nocopy(A.rows(), A.cols(), A.nonZeros(), A.outerIndexPtr(),
+                   A.innerIndexPtr(), A.valuePtr(), A.innerNonZeroPtr());
   Eigen::Map<eig_mat> sparse_map(A.rows(), A.cols(), A.nonZeros(),
-    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
-    A.innerNonZeroPtr());
+                                 A.outerIndexPtr(), A.innerIndexPtr(),
+                                 A.valuePtr(), A.innerNonZeroPtr());
   arena_mat sparse_map_constructor(sparse_map);
   arena_mat A_m(A);
   expect_sparse_matrix_equal(A_m, A);
@@ -148,10 +153,9 @@ TEST(AgradRev, arena_sparse_matrix_constructors) {
   eig_mat mul_B = mul_A * A;
   arena_mat mul_B_m(mul_A_m * A_m);
   expect_sparse_matrix_equal(mul_B_m, mul_B);
-
 }
 
-TEST(AgradRev, arena_sparse_matrix_equals) {
+TEST_F(AgradRev, arena_sparse_matrix_equals) {
   using eig_mat = Eigen::SparseMatrix<double>;
   using stan::math::arena_matrix;
   using arena_mat = arena_matrix<eig_mat>;
@@ -161,22 +165,68 @@ TEST(AgradRev, arena_sparse_matrix_equals) {
   eig_mat A = make_sparse_matrix_random(10, 10);
   // Testing each constructor
   arena_mat empty_m = arena_mat{};
-  arena_mat nocopy = arena_mat(A.rows(), A.cols(), A.nonZeros(),
-    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
-    A.innerNonZeroPtr());
+  empty_m = arena_mat{};
+  arena_mat nocopy
+      = arena_mat(A.rows(), A.cols(), A.nonZeros(), A.outerIndexPtr(),
+                  A.innerIndexPtr(), A.valuePtr(), A.innerNonZeroPtr());
+  nocopy = arena_mat(A.rows(), A.cols(), A.nonZeros(), A.outerIndexPtr(),
+                     A.innerIndexPtr(), A.valuePtr(), A.innerNonZeroPtr());
   Eigen::Map<eig_mat> sparse_map(A.rows(), A.cols(), A.nonZeros(),
-    A.outerIndexPtr(), A.innerIndexPtr(), A.valuePtr(),
-    A.innerNonZeroPtr());
+                                 A.outerIndexPtr(), A.innerIndexPtr(),
+                                 A.valuePtr(), A.innerNonZeroPtr());
   arena_mat sparse_map_constructor = sparse_map;
+  sparse_map_constructor = sparse_map;
   arena_mat A_m = A;
+  A_m = A;
   expect_sparse_matrix_equal(A_m, A);
   arena_mat arena_mat_constructor = A_m;
+  arena_mat_constructor = A_m;
   expect_sparse_matrix_equal(arena_mat_constructor, A);
   eig_mat mul_A = A * A;
   arena_mat mul_A_m = A * A_m;
+  mul_A_m = A * A_m;
   expect_sparse_matrix_equal(mul_A_m, mul_A);
   eig_mat mul_B = mul_A * A;
   arena_mat mul_B_m = mul_A_m * A_m;
+  mul_B_m = mul_A_m * A_m;
   expect_sparse_matrix_equal(mul_B_m, mul_B);
+}
 
+TEST_F(AgradRev, arena_sparse_matrix_inplace_ops) {
+  using eig_mat = Eigen::SparseMatrix<double>;
+  using stan::math::arena_matrix;
+  using arena_mat = arena_matrix<eig_mat>;
+  using inner_iterator = typename eig_mat::InnerIterator;
+  using map_inner_iterator = typename Eigen::Map<eig_mat>::InnerIterator;
+  using stan::test::make_sparse_matrix_random;
+  eig_mat A = make_sparse_matrix_random(10, 10);
+  arena_mat A_m = A;
+  A_m += A;
+  expect_sparse_matrix_equal(A_m, eig_mat(A + A));
+
+  A_m = A;
+  A_m -= A;
+  expect_sparse_matrix_equal(A_m, eig_mat(A - A));
+
+  Eigen::Matrix<double, -1, -1> B
+      = Eigen::Matrix<double, -1, -1>::Random(10, 10);
+  A_m = A;
+  A_m += B;
+  expect_sparse_dense_matrix_equal(A_m, Eigen::Matrix<double, -1, -1>(A + B));
+
+  A_m = A;
+  A_m -= B;
+  expect_sparse_dense_matrix_equal(A_m, Eigen::Matrix<double, -1, -1>(A - B));
+
+  double c = 10;
+  A_m = A;
+  A_m += c;
+  eig_mat C = A;
+  for (int k = 0; k < C.outerSize(); ++k) {
+    inner_iterator it(C, k);
+    for (; bool(it); ++it) {
+      it.valueRef() += c;
+    }
+  }
+  expect_sparse_dense_matrix_equal(A_m, C);
 }
diff --git a/test/unit/math/rev/core/var_test.cpp b/test/unit/math/rev/core/var_test.cpp
index 3e6576e033c..0a56f963453 100644
--- a/test/unit/math/rev/core/var_test.cpp
+++ b/test/unit/math/rev/core/var_test.cpp
@@ -1,6 +1,7 @@
 #include <stan/math.hpp>
 #include <stan/math/prim.hpp>
 #include <test/unit/util.hpp>
+#include <test/unit/math/rev/util.hpp>
 #include <test/unit/pretty_print_types.hpp>
 #include <test/unit/math/rev/fun/util.hpp>
 #include <test/unit/math/rev/core/gradable.hpp>
@@ -8,12 +9,6 @@
 #include <string>
 #include <vector>
 
-struct AgradRev : public testing::Test {
-  void SetUp() {
-    // make sure memory's clean before starting each test
-    stan::math::recover_memory();
-  }
-};
 
 namespace stan {
 namespace test {
@@ -129,7 +124,8 @@ void ctor_overloads_sparse_matrix(EigenMat&& x) {
   inplace_add_var.adj() += test_y;
   // adjoints sparsity pattern will be pattern of x and test_y for addition
   for (int k = 0; k < x.outerSize(); ++k) {
-    for (inner_iterator it(test_y, k), iz(inplace_add_var.adj(), k); iz; ++iz) {
+    typename vari_value<eigen_plain>::InnerIterator iz(inplace_add_var.adj(), k);
+    for (inner_iterator it(test_y, k); bool(iz) && bool(it); ++iz) {
       if (iz.row() == it.row() && iz.col() == it.col()) {
         EXPECT_FLOAT_EQ(iz.value() - 1, it.value());
         ++it;
diff --git a/test/unit/math/rev/core/vari_test.cpp b/test/unit/math/rev/core/vari_test.cpp
index 5aeb921cdc7..80cd80dbe9d 100644
--- a/test/unit/math/rev/core/vari_test.cpp
+++ b/test/unit/math/rev/core/vari_test.cpp
@@ -52,7 +52,8 @@ TEST(AgradRevVari, sparse_matrix_vari) {
   eig_mat B = make_sparse_matrix_random(10, 10);
   vari_value<eig_mat> B_vari(B);
   for (int k = 0; k < B.outerSize(); ++k) {
-    for (inner_iterator it(B, k), iz(B_vari.val(), k); it; ++it, ++iz) {
+    typename vari_value<eig_mat>::InnerIterator iz(B_vari.val(), k);
+    for (inner_iterator it(B, k); bool(it) && bool(iz); ++it, ++iz) {
       EXPECT_FLOAT_EQ(iz.value(), it.value());
     }
   }
diff --git a/test/unit/math/rev/util.hpp b/test/unit/math/rev/util.hpp
index 20b2934561e..87db3580415 100644
--- a/test/unit/math/rev/util.hpp
+++ b/test/unit/math/rev/util.hpp
@@ -5,6 +5,13 @@
 #include <gtest/gtest.h>
 #include <vector>
 
+struct AgradRev : public testing::Test {
+  void SetUp() {
+    // make sure memory's clean before starting each test
+    stan::math::recover_memory();
+  }
+};
+
 namespace stan {
 namespace math {
 namespace test {

From a4939cdccf916877d504ab6a0a62f9f876ac7ccf Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Wed, 8 Nov 2023 17:23:58 -0500
Subject: [PATCH 03/14] fix forward decls for arena_matrix

---
 stan/math/rev/core/arena_matrix.hpp          | 10 +---------
 stan/math/rev/core/var_value_fwd_declare.hpp | 10 ++++++++++
 stan/math/rev/meta/arena_type.hpp            |  7 ++-----
 stan/math/rev/meta/conditional_var_value.hpp |  2 +-
 4 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 61aca79b3c1..0466a7768a0 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -4,20 +4,12 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
 #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>
 #include <iostream>
 namespace stan {
 namespace math {
 
-/**
- * Equivalent to `Eigen::Matrix`, except that the data is stored on AD stack.
- * That makes these objects triviali destructible and usable in `vari`s.
- *
- * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`
- * ...)
- */
-template <typename MatrixType, typename>
-class arena_matrix;
 
 template <typename MatrixType>
 class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>>
diff --git a/stan/math/rev/core/var_value_fwd_declare.hpp b/stan/math/rev/core/var_value_fwd_declare.hpp
index 7a1ef3939ea..d2eef8a9618 100644
--- a/stan/math/rev/core/var_value_fwd_declare.hpp
+++ b/stan/math/rev/core/var_value_fwd_declare.hpp
@@ -6,6 +6,16 @@ namespace math {
 // forward declaration of var
 template <typename T, typename = void>
 class var_value;
+
+/**
+ * Equivalent to `Eigen::Matrix`, except that the data is stored on AD stack.
+ * That makes these objects triviali destructible and usable in `vari`s.
+ *
+ * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`
+ * ...)
+ */
+template <typename MatrixType, typename = void>
+class arena_matrix;
 }  // namespace math
 }  // namespace stan
 #endif
diff --git a/stan/math/rev/meta/arena_type.hpp b/stan/math/rev/meta/arena_type.hpp
index 28c7248b6e5..a7871313146 100644
--- a/stan/math/rev/meta/arena_type.hpp
+++ b/stan/math/rev/meta/arena_type.hpp
@@ -6,13 +6,10 @@
 #include <stan/math/prim/meta/plain_type.hpp>
 #include <stan/math/rev/core/arena_allocator.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
+#include <stan/math/rev/core/var_value_fwd_declare.hpp>
+
 
 namespace stan {
-namespace math {
-// forward declaration
-template <typename MatrixType, typename = void>
-class arena_matrix;
-}  // namespace math
 
 namespace internal {
 template <typename T, typename = void, typename = void>
diff --git a/stan/math/rev/meta/conditional_var_value.hpp b/stan/math/rev/meta/conditional_var_value.hpp
index 3ad0f2604f8..cc9a54c9406 100644
--- a/stan/math/rev/meta/conditional_var_value.hpp
+++ b/stan/math/rev/meta/conditional_var_value.hpp
@@ -1,7 +1,7 @@
 #ifndef STAN_MATH_REV_META_CONDITIONAL_VAR_EIGEN_HPP
 #define STAN_MATH_REV_META_CONDITIONAL_VAR_EIGEN_HPP
 
-#include <stan/math/rev/core/var.hpp>
+#include <stan/math/rev/core/var_value_fwd_declare.hpp>
 #include <stan/math/rev/meta/plain_type.hpp>
 #include <stan/math/prim/meta/promote_scalar_type.hpp>
 

From 3889b48f853ac2f0fda9e4b20c8443f06aa023e5 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 13 Nov 2023 09:54:50 -0500
Subject: [PATCH 04/14] remove iostream header from arena_matrix.hpp

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

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 0466a7768a0..d9912eb9adc 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -6,7 +6,7 @@
 #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>
-#include <iostream>
+
 namespace stan {
 namespace math {
 

From 3ec50c792904b24cfd908e1ac79abf4ecb97a172 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 20 Nov 2023 12:36:06 -0500
Subject: [PATCH 05/14] fix cpplint

---
 stan/math/rev/core/arena_matrix.hpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index d9912eb9adc..d9a11dd1ff1 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -287,7 +287,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
     for (int k = 0; k < (*this).outerSize(); ++k) {
       typename Base::InnerIterator it(*this, k);
       typename std::decay_t<Expr>::InnerIterator iz(x, k);
-      for (; bool(it) && bool(iz); ++it, ++iz) {
+      for (; static_cast<bool>(it) && static_cast<bool>(iz); ++it, ++iz) {
         f(it.valueRef(), iz.value());
       }
     }
@@ -310,7 +310,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
     auto&& x = to_ref(other);
     for (int k = 0; k < (*this).outerSize(); ++k) {
       typename Base::InnerIterator it(*this, k);
-      for (; bool(it); ++it) {
+      for (; static_cast<bool>(it); ++it) {
         f(it.valueRef(), x(it.row(), it.col()));
       }
     }
@@ -331,7 +331,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
   void inplace_ops_impl(F&& f, T&& other) {
     for (int k = 0; k < (*this).outerSize(); ++k) {
       typename Base::InnerIterator it(*this, k);
-      for (; bool(it); ++it) {
+      for (; static_cast<bool>(it); ++it) {
         f(it.valueRef(), other);
       }
     }

From cc77b65a1bc50411c2ee16c86f476c43ed20e90d Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Mon, 20 Nov 2023 12:49:55 -0500
Subject: [PATCH 06/14] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 stan/math/rev/core/arena_matrix.hpp  | 1 -
 stan/math/rev/core/vari.hpp          | 3 +--
 stan/math/rev/meta/arena_type.hpp    | 1 -
 test/unit/math/rev/core/var_test.cpp | 4 ++--
 4 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index d9a11dd1ff1..069228a40f1 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -10,7 +10,6 @@
 namespace stan {
 namespace math {
 
-
 template <typename MatrixType>
 class arena_matrix<MatrixType, require_eigen_dense_base_t<MatrixType>>
     : public Eigen::Map<MatrixType> {
diff --git a/stan/math/rev/core/vari.hpp b/stan/math/rev/core/vari.hpp
index f3caf6f925a..80da3b0f46d 100644
--- a/stan/math/rev/core/vari.hpp
+++ b/stan/math/rev/core/vari.hpp
@@ -869,8 +869,7 @@ 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)) {
+  vari_value(S&& x, bool stacked) : adj_(x), val_(std::forward<S>(x)) {
     this->set_zero_adjoint();
     if (stacked) {
       ChainableStack::instance_->var_stack_.push_back(this);
diff --git a/stan/math/rev/meta/arena_type.hpp b/stan/math/rev/meta/arena_type.hpp
index a7871313146..095fd53f86f 100644
--- a/stan/math/rev/meta/arena_type.hpp
+++ b/stan/math/rev/meta/arena_type.hpp
@@ -8,7 +8,6 @@
 #include <stan/math/rev/core/chainable_alloc.hpp>
 #include <stan/math/rev/core/var_value_fwd_declare.hpp>
 
-
 namespace stan {
 
 namespace internal {
diff --git a/test/unit/math/rev/core/var_test.cpp b/test/unit/math/rev/core/var_test.cpp
index 0a56f963453..341302b1093 100644
--- a/test/unit/math/rev/core/var_test.cpp
+++ b/test/unit/math/rev/core/var_test.cpp
@@ -9,7 +9,6 @@
 #include <string>
 #include <vector>
 
-
 namespace stan {
 namespace test {
 template <typename T, typename S>
@@ -124,7 +123,8 @@ void ctor_overloads_sparse_matrix(EigenMat&& x) {
   inplace_add_var.adj() += test_y;
   // adjoints sparsity pattern will be pattern of x and test_y for addition
   for (int k = 0; k < x.outerSize(); ++k) {
-    typename vari_value<eigen_plain>::InnerIterator iz(inplace_add_var.adj(), k);
+    typename vari_value<eigen_plain>::InnerIterator iz(inplace_add_var.adj(),
+                                                       k);
     for (inner_iterator it(test_y, k); bool(iz) && bool(it); ++iz) {
       if (iz.row() == it.row() && iz.col() == it.col()) {
         EXPECT_FLOAT_EQ(iz.value() - 1, it.value());

From ed6877f56d2de8d3e627640594217cdc09a9af76 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 20 Nov 2023 12:58:18 -0500
Subject: [PATCH 07/14] update

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

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 069228a40f1..56a3175f582 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -339,7 +339,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
  public:
   /**
    * Inplace operator `+=`
-   * @note Caution!! Inplace operators assume that either
+   * @note Caution! Inplace operators assume that either
    *  1. The right hand side sparse matrix has the same sparcity pattern
    *  2. You only intend to add a scalar or dense matrix coefficients to the
    * nonzero values of `this` This is intended to be used within the reverse

From ebe65ad052e53ff73026ccb7a50c94444be2603f Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Mon, 20 Nov 2023 14:18:28 -0500
Subject: [PATCH 08/14] update

---
 test/unit/math/rev/core/vari_test.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/unit/math/rev/core/vari_test.cpp b/test/unit/math/rev/core/vari_test.cpp
index 80cd80dbe9d..a543941c964 100644
--- a/test/unit/math/rev/core/vari_test.cpp
+++ b/test/unit/math/rev/core/vari_test.cpp
@@ -53,7 +53,7 @@ TEST(AgradRevVari, sparse_matrix_vari) {
   vari_value<eig_mat> B_vari(B);
   for (int k = 0; k < B.outerSize(); ++k) {
     typename vari_value<eig_mat>::InnerIterator iz(B_vari.val(), k);
-    for (inner_iterator it(B, k); bool(it) && bool(iz); ++it, ++iz) {
+    for (inner_iterator it(B, k); static_cast<bool>(it) && static_cast<bool>(iz); ++it, ++iz) {
       EXPECT_FLOAT_EQ(iz.value(), it.value());
     }
   }

From 700d31a393ec9aa25d122e0ba643365d265ab4c6 Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Mon, 20 Nov 2023 14:19:26 -0500
Subject: [PATCH 09/14] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 test/unit/math/rev/core/vari_test.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/unit/math/rev/core/vari_test.cpp b/test/unit/math/rev/core/vari_test.cpp
index a543941c964..3f0ad19d5c5 100644
--- a/test/unit/math/rev/core/vari_test.cpp
+++ b/test/unit/math/rev/core/vari_test.cpp
@@ -53,7 +53,8 @@ TEST(AgradRevVari, sparse_matrix_vari) {
   vari_value<eig_mat> B_vari(B);
   for (int k = 0; k < B.outerSize(); ++k) {
     typename vari_value<eig_mat>::InnerIterator iz(B_vari.val(), k);
-    for (inner_iterator it(B, k); static_cast<bool>(it) && static_cast<bool>(iz); ++it, ++iz) {
+    for (inner_iterator it(B, k);
+         static_cast<bool>(it) && static_cast<bool>(iz); ++it, ++iz) {
       EXPECT_FLOAT_EQ(iz.value(), it.value());
     }
   }

From e354cc0d87d444392cded57e3496dbff35f41752 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Fri, 1 Dec 2023 15:04:57 -0500
Subject: [PATCH 10/14] update for cpplint

---
 test/unit/math/rev/core/arena_matrix_test.cpp | 2 +-
 test/unit/math/rev/core/var_test.cpp          | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/unit/math/rev/core/arena_matrix_test.cpp b/test/unit/math/rev/core/arena_matrix_test.cpp
index 5ed3ac0a6c2..3e1bfe4dc44 100644
--- a/test/unit/math/rev/core/arena_matrix_test.cpp
+++ b/test/unit/math/rev/core/arena_matrix_test.cpp
@@ -224,7 +224,7 @@ TEST_F(AgradRev, arena_sparse_matrix_inplace_ops) {
   eig_mat C = A;
   for (int k = 0; k < C.outerSize(); ++k) {
     inner_iterator it(C, k);
-    for (; bool(it); ++it) {
+    for (; static_cast<bool>(it); ++it) {
       it.valueRef() += c;
     }
   }
diff --git a/test/unit/math/rev/core/var_test.cpp b/test/unit/math/rev/core/var_test.cpp
index 341302b1093..b83a9360fdc 100644
--- a/test/unit/math/rev/core/var_test.cpp
+++ b/test/unit/math/rev/core/var_test.cpp
@@ -125,7 +125,7 @@ void ctor_overloads_sparse_matrix(EigenMat&& x) {
   for (int k = 0; k < x.outerSize(); ++k) {
     typename vari_value<eigen_plain>::InnerIterator iz(inplace_add_var.adj(),
                                                        k);
-    for (inner_iterator it(test_y, k); bool(iz) && bool(it); ++iz) {
+    for (inner_iterator it(test_y, k); static_cast<bool>(iz) && static_cast<bool>(it); ++iz) {
       if (iz.row() == it.row() && iz.col() == it.col()) {
         EXPECT_FLOAT_EQ(iz.value() - 1, it.value());
         ++it;

From 8fab167c68fe78eea0ab49f41fe8559387963995 Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Fri, 1 Dec 2023 15:06:07 -0500
Subject: [PATCH 11/14] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 test/unit/math/rev/core/var_test.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/unit/math/rev/core/var_test.cpp b/test/unit/math/rev/core/var_test.cpp
index 89fb06caeba..a2ca9a36d64 100644
--- a/test/unit/math/rev/core/var_test.cpp
+++ b/test/unit/math/rev/core/var_test.cpp
@@ -125,7 +125,8 @@ void ctor_overloads_sparse_matrix(EigenMat&& x) {
   for (int k = 0; k < x.outerSize(); ++k) {
     typename vari_value<eigen_plain>::InnerIterator iz(inplace_add_var.adj(),
                                                        k);
-    for (inner_iterator it(test_y, k); static_cast<bool>(iz) && static_cast<bool>(it); ++iz) {
+    for (inner_iterator it(test_y, k);
+         static_cast<bool>(iz) && static_cast<bool>(it); ++iz) {
       if (iz.row() == it.row() && iz.col() == it.col()) {
         EXPECT_FLOAT_EQ(iz.value() - 1, it.value());
         ++it;

From 0656caa06dda69d685dccbcd6db4336cf498572d Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Thu, 29 Feb 2024 11:58:19 -0500
Subject: [PATCH 12/14] update docs

---
 stan/math/rev/core/arena_matrix.hpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 24832b90466..182972ae197 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -200,8 +200,8 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   template <typename S, require_convertible_t<S&, PlainObject>* = nullptr,
             require_not_same_t<S, PlainObject>* = nullptr>
-  arena_matrix(S&& x)  // NOLINT
-      : arena_matrix(PlainObject(x)) {}
+  arena_matrix(S&& other)  // NOLINT
+      : arena_matrix(PlainObject(std::forward<S>(other))) {}
 
   /**
    * Constructs `arena_matrix` from an expression. This makes an assumption that

From dfd3949c17af8fe684ba19b0e7a69f01ddc29bf5 Mon Sep 17 00:00:00 2001
From: Steve Bronder <sbronder@flatironinstitute.org>
Date: Tue, 26 Mar 2024 11:52:14 -0400
Subject: [PATCH 13/14] cleanup docs, add innerNonZero() calls, speedup inplace
 ops for sparse matrix

---
 stan/math/rev/core/arena_matrix.hpp           | 74 +++++++++++++------
 stan/math/rev/core/var_value_fwd_declare.hpp  |  4 +-
 test/unit/math/rev/core/arena_matrix_test.cpp | 11 ++-
 3 files changed, 62 insertions(+), 27 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 182972ae197..88215ed9e7f 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -176,7 +176,8 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
 
  private:
   template <typename T, typename Integral>
-  auto* copy_vector(const T* ptr, Integral size) {
+  inline T* copy_vector(const T* ptr, Integral size) {
+    if (size == 0) return nullptr;
     T* ret = ChainableStack::instance_->memalloc_.alloc_array<T>(size);
     std::copy_n(ptr, size, ret);
     return ret;
@@ -192,7 +193,8 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
       : Base::Map(other.rows(), other.cols(), other.nonZeros(),
                   copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
                   copy_vector(other.innerIndexPtr(), other.nonZeros()),
-                  copy_vector(other.valuePtr(), other.nonZeros())) {}
+                  copy_vector(other.valuePtr(), other.nonZeros()),
+                  copy_vector(other.innerNonZeroPtr(), other.innerNonZeroPtr() == nullptr ? 0 : other.innerSize())) {}
 
   /**
    * Constructs `arena_matrix` from an Eigen expression
@@ -220,7 +222,7 @@ 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.valuePtr(), other.innernonZeroPtr()) {}
   /**
    * Move constructor.
    * @note Since the memory for the arena matrix sits in Stan's memory arena all
@@ -230,7 +232,7 @@ 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.valuePtr(), other.innerNonZeroPtr()) {}
   /**
    * Copy constructor. No actual copy is performed
    * @note Since the memory for the arena matrix sits in Stan's memory arena all
@@ -240,7 +242,7 @@ 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.valuePtr(), other.innerNonZeroPtr()) {}
 
   // without this using, compiler prefers combination of implicit construction
   // and copy assignment to the inherited operator when assigned an expression
@@ -259,7 +261,8 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
     new (this) Base(other.rows(), other.cols(), other.nonZeros(),
                     const_cast<StorageIndex*>(other.outerIndexPtr()),
                     const_cast<StorageIndex*>(other.innerIndexPtr()),
-                    const_cast<Scalar*>(other.valuePtr()));
+                    const_cast<Scalar*>(other.valuePtr()),
+                    const_cast<StorageIndex*>(other.innerNonZeroPtr()));
     return *this;
   }
 
@@ -289,8 +292,32 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    * @param other The right hand side of the inplace operation
    */
   template <typename F, typename Expr,
-            require_convertible_t<Expr&, MatrixType>* = nullptr>
-  void inplace_ops_impl(F&& f, Expr&& other) {
+            require_convertible_t<Expr&, MatrixType>* = nullptr,
+            require_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  inline void inplace_ops_impl(F&& f, Expr&& other) {
+    auto&& x = to_ref(other);
+    auto* val_ptr = (*this).valuePtr();
+    auto* x_val_ptr = x.valuePtr();
+    const auto non_zeros = (*this).nonZeros();
+    for (Eigen::Index i = 0; i < non_zeros; ++i) {
+      f(val_ptr[i], x_val_ptr[i]);
+    }
+  }
+
+  /**
+   * inplace operations functor for `Sparse (.*)= Sparse`.
+   * @note This assumes that each matrix is of the same size and sparsity
+   * pattern.
+   * @tparam F A type with a valid `operator()(Scalar& x, const Scalar& y)`
+   * method
+   * @tparam Expr Type derived from `Eigen::SparseMatrixBase`
+   * @param f Functor that performs the inplace operation
+   * @param other The right hand side of the inplace operation
+   */
+  template <typename F, typename Expr,
+            require_convertible_t<Expr&, MatrixType>* = nullptr,
+            require_not_same_t<Expr, arena_matrix<MatrixType>>* = nullptr>
+  inline void inplace_ops_impl(F&& f, Expr&& other) {
     auto&& x = to_ref(other);
     for (int k = 0; k < (*this).outerSize(); ++k) {
       typename Base::InnerIterator it(*this, k);
@@ -314,7 +341,7 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
   template <typename F, typename Expr,
             require_not_convertible_t<Expr&, MatrixType>* = nullptr,
             require_eigen_dense_base_t<Expr>* = nullptr>
-  void inplace_ops_impl(F&& f, Expr&& other) {
+  inline void inplace_ops_impl(F&& f, Expr&& other) {
     auto&& x = to_ref(other);
     for (int k = 0; k < (*this).outerSize(); ++k) {
       typename Base::InnerIterator it(*this, k);
@@ -336,12 +363,11 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   template <typename F, typename T,
             require_convertible_t<T&, Scalar>* = nullptr>
-  void inplace_ops_impl(F&& f, T&& other) {
-    for (int k = 0; k < (*this).outerSize(); ++k) {
-      typename Base::InnerIterator it(*this, k);
-      for (; static_cast<bool>(it); ++it) {
-        f(it.valueRef(), other);
-      }
+  inline void inplace_ops_impl(F&& f, T&& other) {
+    auto* val_ptr = (*this).valuePtr();
+    const auto non_zeros = (*this).nonZeros();
+    for (Eigen::Index i = 0; i < non_zeros; ++i) {
+      f(val_ptr[i], other);
     }
   }
 
@@ -349,17 +375,17 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
   /**
    * Inplace operator `+=`
    * @note Caution! Inplace operators assume that either
-   *  1. The right hand side sparse matrix has the same sparcity pattern
+   *  1. The right hand side sparse matrix has the same sparsity pattern
    *  2. You only intend to add a scalar or dense matrix coefficients to the
-   * nonzero values of `this` This is intended to be used within the reverse
+   * nonzero values of `this`. This is intended to be used within the reverse
    * pass for accumulation of the adjoint and is built as such. Any other use
-   * case should be be sure the above assumptions are satisfied.
+   * case should be be sure that the above assumptions are satisfied.
    * @tparam T A type derived from `Eigen::SparseMatrixBase` or
    * `Eigen::DenseMatrixBase` or a `Scalar`
    * @param other value to be added inplace to the matrix.
    */
   template <typename T>
-  arena_matrix& operator+=(T&& other) {
+  inline arena_matrix& operator+=(T&& other) {
     inplace_ops_impl(
         [](auto&& x, auto&& y) {
           x += y;
@@ -371,18 +397,18 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
 
   /**
    * Inplace operator `+=`
-   * @note Caution!! Inplace operators assume that either
-   *  1. The right hand side sparse matrix has the same sparcity pattern
+   * @note Caution! Inplace operators assume that either
+   *  1. The right hand side sparse matrix has the same sparsity pattern
    *  2. You only intend to add a scalar or dense matrix coefficients to the
-   * nonzero values of `this` This is intended to be used within the reverse
+   * nonzero values of `this`. This is intended to be used within the reverse
    * pass for accumulation of the adjoint and is built as such. Any other use
-   * case should be be sure the above assumptions are satisfied.
+   * case should be be sure that the above assumptions are satisfied.
    * @tparam T A type derived from `Eigen::SparseMatrixBase` or
    * `Eigen::DenseMatrixBase` or a `Scalar`
    * @param other value to be added inplace to the matrix.
    */
   template <typename T>
-  arena_matrix& operator-=(T&& other) {
+  inline arena_matrix& operator-=(T&& other) {
     inplace_ops_impl(
         [](auto&& x, auto&& y) {
           x -= y;
diff --git a/stan/math/rev/core/var_value_fwd_declare.hpp b/stan/math/rev/core/var_value_fwd_declare.hpp
index d2eef8a9618..15779b3ed36 100644
--- a/stan/math/rev/core/var_value_fwd_declare.hpp
+++ b/stan/math/rev/core/var_value_fwd_declare.hpp
@@ -9,9 +9,9 @@ class var_value;
 
 /**
  * Equivalent to `Eigen::Matrix`, except that the data is stored on AD stack.
- * That makes these objects triviali destructible and usable in `vari`s.
+ * That makes these objects trivially destructible and usable in `vari`s.
  *
- * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`
+ * @tparam MatrixType Eigen matrix type this works as (`MatrixXd`, `VectorXd`,
  * ...)
  */
 template <typename MatrixType, typename = void>
diff --git a/test/unit/math/rev/core/arena_matrix_test.cpp b/test/unit/math/rev/core/arena_matrix_test.cpp
index 3e1bfe4dc44..cb339ae73bd 100644
--- a/test/unit/math/rev/core/arena_matrix_test.cpp
+++ b/test/unit/math/rev/core/arena_matrix_test.cpp
@@ -136,7 +136,7 @@ TEST_F(AgradRev, arena_sparse_matrix_constructors) {
   using stan::test::make_sparse_matrix_random;
   eig_mat A = make_sparse_matrix_random(10, 10);
   // Testing each constructor
-  arena_mat empty_m();
+  arena_mat empty_m{};
   arena_mat nocopy(A.rows(), A.cols(), A.nonZeros(), A.outerIndexPtr(),
                    A.innerIndexPtr(), A.valuePtr(), A.innerNonZeroPtr());
   Eigen::Map<eig_mat> sparse_map(A.rows(), A.cols(), A.nonZeros(),
@@ -208,6 +208,15 @@ TEST_F(AgradRev, arena_sparse_matrix_inplace_ops) {
   A_m -= A;
   expect_sparse_matrix_equal(A_m, eig_mat(A - A));
 
+  arena_mat B_m = A;
+  A_m = A;
+  A_m += B_m;
+  expect_sparse_matrix_equal(A_m, eig_mat(A + A));
+
+  A_m = A;
+  A_m -= B_m;
+  expect_sparse_matrix_equal(A_m, eig_mat(A - A));
+
   Eigen::Matrix<double, -1, -1> B
       = Eigen::Matrix<double, -1, -1>::Random(10, 10);
   A_m = A;

From 11e414ec9abfc4afb721a6241ab98017edfbde4c Mon Sep 17 00:00:00 2001
From: Stan Jenkins <mc.stanislaw@gmail.com>
Date: Tue, 26 Mar 2024 11:53:12 -0400
Subject: [PATCH 14/14] [Jenkins] auto-formatting by clang-format version
 10.0.0-4ubuntu1

---
 stan/math/rev/core/arena_matrix.hpp | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/stan/math/rev/core/arena_matrix.hpp b/stan/math/rev/core/arena_matrix.hpp
index 88215ed9e7f..69baf5b81c0 100644
--- a/stan/math/rev/core/arena_matrix.hpp
+++ b/stan/math/rev/core/arena_matrix.hpp
@@ -177,7 +177,8 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
  private:
   template <typename T, typename Integral>
   inline T* copy_vector(const T* ptr, Integral size) {
-    if (size == 0) return nullptr;
+    if (size == 0)
+      return nullptr;
     T* ret = ChainableStack::instance_->memalloc_.alloc_array<T>(size);
     std::copy_n(ptr, size, ret);
     return ret;
@@ -190,11 +191,14 @@ class arena_matrix<MatrixType, require_eigen_sparse_base_t<MatrixType>>
    */
   template <typename T, require_same_t<T, PlainObject>* = nullptr>
   arena_matrix(T&& other)  // NOLINT
-      : Base::Map(other.rows(), other.cols(), other.nonZeros(),
-                  copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
-                  copy_vector(other.innerIndexPtr(), other.nonZeros()),
-                  copy_vector(other.valuePtr(), other.nonZeros()),
-                  copy_vector(other.innerNonZeroPtr(), other.innerNonZeroPtr() == nullptr ? 0 : other.innerSize())) {}
+      : Base::Map(
+          other.rows(), other.cols(), other.nonZeros(),
+          copy_vector(other.outerIndexPtr(), other.outerSize() + 1),
+          copy_vector(other.innerIndexPtr(), other.nonZeros()),
+          copy_vector(other.valuePtr(), other.nonZeros()),
+          copy_vector(
+              other.innerNonZeroPtr(),
+              other.innerNonZeroPtr() == nullptr ? 0 : other.innerSize())) {}
 
   /**
    * Constructs `arena_matrix` from an Eigen expression