Vector-valued functions with manual gradients in external C++

General
,
3

Bob’s answer is right that make_callback_var() is meant to return stan::math::var_value<T> types. In additon to his fix I would make an explicit version for var return types like the following

template<typename T1__, stan::require_all_t<stan::is_col_vector<T1__>,
                                            stan::is_vt_not_complex<T1__>, stan::is_vt_var<T1__>>* = nullptr>
Eigen::Matrix<stan::math::var, -1, 1>

odefunc(const double& t, const T1__& y, const double& theta,
        std::ostream* pstream__)
{
  Eigen::Matrix<double, -1, 1> dydt =
    Eigen::Matrix<double, -1, 1>::Constant(2, 0.0);
  dydt[0] = y[1].val();
  dydt[1] = -theta * y[0].val();

  stan::arena_t<Eigen::Matrix<double, -1, 1>> jacobian = Eigen::Matrix<double, -1, -1>::Constant(2, 2, 0.0);
  jacobian[0, 1] = 1.0;
  jacobian[1, 0] = -theta;

  stan::arena_t<T1__> y_arena = y;
  // Everything matrix/vector that is used in the reverse pass must be put on the arena
  reverse_pass_callback([y_arena, jacobian]() mutable {
    y_arena.adj() += jacobian.transpose() * y_arena.adj();
  });
  return Eigen::Matrix<stan::math::var, -1, 1>(dydt);
}

Though I think having a specialization for var means the compiler may also ask about a version for fvar<T>. If the compiler doesn’t complain then dont worry about it, but the good news is that if it does complain then you can just define the forward mode specialization and leave it empty since the ode solvers never call forward mode autodiff.