update multiply for single to double loop

Author: SteveBronder

stan/math/rev/fun/multiply.hpp

@@ -135,31 +135,37 @@ template <typename T1, typename T2, require_not_matrix_t<T1>* = nullptr,
           require_not_row_and_col_vector_t<T1, T2>* = nullptr>
 inline auto multiply(const T1& A, const T2& B) {
   if (!is_constant<T2>::value && !is_constant<T1>::value) {
-    arena_t<promote_scalar_t<var, T2>> arena_B = to_ref(B);
-    arena_t<promote_scalar_t<double, T2>> arena_B_val = value_of(arena_B);
+    arena_t<promote_scalar_t<var, T1>> arena_A = A;
+    arena_t<promote_scalar_t<var, T2>> arena_B = B;
     using return_t = return_var_matrix_t<T2, T1, T2>;
-    arena_t<return_t> res = value_of(A) * arena_B_val.array();
-    reverse_pass_callback([A, arena_B, arena_B_val, res]() mutable {
-      auto res_adj = res.adj().eval();
-      forward_as<var>(A).adj() += (res_adj.array() * arena_B_val.array()).sum();
-      arena_B.adj().array() += value_of(A) * res_adj.array();
+    arena_t<return_t> res = arena_A.val() * arena_B.val().array();
+    reverse_pass_callback([arena_A, arena_B, res]() mutable {
+      const auto a_val = arena_A.val();
+      for (Eigen::Index j = 0; j < res.cols(); ++j) {
+        for (Eigen::Index i = 0; i < res.rows(); ++i) {
+          const auto res_adj = res.adj().coeffRef(i, j);
+          arena_A.adj() += res_adj * arena_B.val().coeff(i, j);
+          arena_B.adj().coeffRef(i, j) += a_val * res_adj;
+        }
+      }
     });
     return return_t(res);
   } else if (!is_constant<T2>::value) {
-    arena_t<promote_scalar_t<var, T2>> arena_B = to_ref(B);
+    arena_t<promote_scalar_t<double, T1>> arena_A = value_of(A);
+    arena_t<promote_scalar_t<var, T2>> arena_B = B;
     using return_t = return_var_matrix_t<T2, T1, T2>;
-    arena_t<return_t> res = value_of(A) * value_of(arena_B).array();
-    reverse_pass_callback([A, arena_B, res]() mutable {
-      arena_B.adj().array() += value_of(A) * res.adj().array();
+    arena_t<return_t> res = arena_A * arena_B.val().array();
+    reverse_pass_callback([arena_A, arena_B, res]() mutable {
+      arena_B.adj().array() += arena_A * res.adj().array();
     });
     return return_t(res);
   } else {
-    arena_t<promote_scalar_t<double, T2>> arena_B_val = value_of(B);
+    arena_t<promote_scalar_t<var, T1>> arena_A = A;
+    arena_t<promote_scalar_t<double, T2>> arena_B = value_of(B);
     using return_t = return_var_matrix_t<T2, T1, T2>;
-    arena_t<return_t> res = value_of(A) * arena_B_val.array();
-    reverse_pass_callback([A, arena_B_val, res]() mutable {
-      forward_as<var>(A).adj()
-          += (res.adj().array() * arena_B_val.array()).sum();
+    arena_t<return_t> res = arena_A.val() * arena_B.array();
+    reverse_pass_callback([arena_A, arena_B, res]() mutable {
+      arena_A.adj() += (res.adj().array() * arena_B.array()).sum();
     });
     return return_t(res);
   }