Feature/adjoint odes

stan/math/rev/core/zero_adjoints.hpp

@@ -8,24 +8,10 @@
 namespace stan {
 namespace math {
 
-inline void zero_adjoints();
-
-template <typename T, typename... Pargs, require_st_arithmetic<T>* = nullptr>
-inline void zero_adjoints(T& x, Pargs&... args);
-
-template <typename... Pargs>
-inline void zero_adjoints(var& x, Pargs&... args);
-
-template <int R, int C, typename... Pargs>
-inline void zero_adjoints(Eigen::Matrix<var, R, C>& x, Pargs&... args);
-
-template <typename T, typename... Pargs, require_st_autodiff<T>* = nullptr>
-inline void zero_adjoints(std::vector<T>& x, Pargs&... args);
-
 /**
  * End of recursion for set_zero_adjoints
  */
-inline void zero_adjoints() {}
+inline void zero_adjoints() noexcept {}
 
 /**
  * Do nothing for non-autodiff arguments. Recursively call zero_adjoints
@@ -37,10 +23,8 @@ inline void zero_adjoints() {}
  * @param x current argument
  * @param args rest of arguments to zero
  */
-template <typename T, typename... Pargs, require_st_arithmetic<T>*>
-inline void zero_adjoints(T& x, Pargs&... args) {
-  zero_adjoints(args...);
-}
+template <typename T, require_st_arithmetic<T>* = nullptr>
+inline void zero_adjoints(T& x) noexcept {}
 
 /**
  * Zero the adjoint of the vari in the first argument. Recursively call
@@ -52,11 +36,7 @@ inline void zero_adjoints(T& x, Pargs&... args) {
  * @param x current argument
  * @param args rest of arguments to zero
  */
-template <typename... Pargs>
-inline void zero_adjoints(var& x, Pargs&... args) {
-  x.vi_->set_zero_adjoint();
-  zero_adjoints(args...);
-}
+inline void zero_adjoints(var& x) { x.adj() = 0; }
 
 /**
  * Zero the adjoints of the varis of every var in an Eigen::Matrix
@@ -68,11 +48,10 @@ inline void zero_adjoints(var& x, Pargs&... args) {
  * @param x current argument
  * @param args rest of arguments to zero
  */
-template <int R, int C, typename... Pargs>
-inline void zero_adjoints(Eigen::Matrix<var, R, C>& x, Pargs&... args) {
+template <typename EigMat, require_eigen_vt<is_autodiff, EigMat>* = nullptr>
+inline void zero_adjoints(EigMat& x) {
   for (size_t i = 0; i < x.size(); ++i)
-    x.coeffRef(i).vi_->set_zero_adjoint();
-  zero_adjoints(args...);
+    x.coeffRef(i).adj() = 0;
 }
 
 /**
@@ -85,11 +64,12 @@ inline void zero_adjoints(Eigen::Matrix<var, R, C>& x, Pargs&... args) {
  * @param x current argument
  * @param args rest of arguments to zero
  */
-template <typename T, typename... Pargs, require_st_autodiff<T>*>
-inline void zero_adjoints(std::vector<T>& x, Pargs&... args) {
-  for (size_t i = 0; i < x.size(); ++i)
+template <typename StdVec,
+          require_std_vector_st<is_autodiff, StdVec>* = nullptr>
+inline void zero_adjoints(StdVec& x) {
+  for (size_t i = 0; i < x.size(); ++i) {
     zero_adjoints(x[i]);
-  zero_adjoints(args...);
+  }
 }
 
 }  // namespace math

stan/math/rev/functor.hpp

@@ -17,6 +17,7 @@
 #include <stan/math/rev/functor/integrate_ode_bdf.hpp>
 #include <stan/math/rev/functor/ode_adams.hpp>
 #include <stan/math/rev/functor/ode_bdf.hpp>
+#include <stan/math/rev/functor/ode_adjoint.hpp>
 #include <stan/math/rev/functor/ode_store_sensitivities.hpp>
 #include <stan/math/rev/functor/jacobian.hpp>
 #include <stan/math/rev/functor/kinsol_data.hpp>

stan/math/rev/functor/coupled_ode_system.hpp

@@ -6,6 +6,7 @@
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/rev/fun/value_of.hpp>
+#include <stan/math/prim/functor/for_each.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stdexcept>
 #include <ostream>
@@ -137,8 +138,10 @@ struct coupled_ode_system_impl<false, F, T_y0, Args...> {
 
       y_adjoints_ = y_vars.adj();
 
-      // memset was faster than Eigen setZero
-      memset(args_adjoints_.data(), 0, sizeof(double) * num_args_vars);
+      if (args_adjoints_.size() > 0) {
+        memset(args_adjoints_.data(), 0,
+               sizeof(double) * args_adjoints_.size());
+      }
 
       apply(
           [&](auto&&... args) {
@@ -148,7 +151,8 @@ struct coupled_ode_system_impl<false, F, T_y0, Args...> {
 
       // The vars here do not live on the nested stack so must be zero'd
       // separately
-      apply([&](auto&&... args) { zero_adjoints(args...); }, local_args_tuple_);
+      stan::math::for_each([](auto&& arg) { zero_adjoints(arg); },
+                           local_args_tuple_);
 
       // No need to zero adjoints after last sweep
       if (i + 1 < N_) {