1 #ifndef STAN_MATH_REV_MAT_FUNCTOR_ADJ_JAC_APPLY_HPP 2 #define STAN_MATH_REV_MAT_FUNCTOR_ADJ_JAC_APPLY_HPP 28 template <
class F,
class Tuple, std::size_t... I>
30 std::index_sequence<I...> i) {
31 return f(std::get<I>(t)...);
46 template <
class F,
class Tuple>
47 constexpr
auto apply(
const F& f,
const Tuple& t) {
48 return apply_impl(f, t, std::make_index_sequence<std::tuple_size<Tuple>{}>{});
59 template <
size_t size>
61 y_adj = y_vi[0]->
adj_;
72 template <
size_t size>
74 std::vector<double>& y_adj) {
76 for (
size_t m = 0; m < y_adj.size(); ++m)
77 y_adj[m] = y_vi[m]->adj_;
88 template <
size_t size,
int R,
int C>
90 Eigen::Matrix<double, R, C>& y_adj) {
91 y_adj.resize(M[0], M[1]);
92 for (
int m = 0; m < y_adj.size(); ++m)
93 y_adj(m) = y_vi[m]->
adj_;
101 template <
typename T>
110 static constexpr
size_t value = 0;
117 template <
typename T>
119 static constexpr
size_t value = 1;
126 template <
typename T,
int R,
int C>
128 static constexpr
size_t value = 2;
146 template <
typename F,
typename... Targs>
148 std::array<bool,
sizeof...(Targs)> is_var_;
150 = std::result_of_t<F(decltype(is_var_), decltype(
value_of(Targs()))...)>;
153 std::array<int,
sizeof...(Targs)> offsets_;
155 std::array<int, internal::compute_dims<FReturnType>::value>
M_;
181 template <
int R,
int C,
typename... Pargs>
183 const Pargs&... args) {
184 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
187 return count_memory(count, args...);
190 template <
int R,
int C,
typename... Pargs>
191 size_t count_memory(
size_t count,
const Eigen::Matrix<double, R, C>& x,
192 const Pargs&... args) {
193 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
195 return count_memory(count, args...);
198 template <
typename... Pargs>
200 const Pargs&... args) {
201 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
204 return count_memory(count, args...);
207 template <
typename... Pargs>
209 const Pargs&... args) {
210 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
212 return count_memory(count, args...);
215 template <
typename... Pargs>
217 const Pargs&... args) {
218 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
220 return count_memory(count, args...);
223 template <
typename... Pargs>
225 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
228 return count_memory(count, args...);
231 template <
typename... Pargs>
232 size_t count_memory(
size_t count,
const double& x,
const Pargs&... args) {
233 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
235 return count_memory(count, args...);
238 template <
typename... Pargs>
239 size_t count_memory(
size_t count,
const int& x,
const Pargs&... args) {
240 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
242 return count_memory(count, args...);
264 template <
int R,
int C,
typename... Pargs>
265 void prepare_x_vis(
const Eigen::Matrix<var, R, C>& x,
const Pargs&... args) {
266 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
267 for (
int i = 0; i < x.size(); ++i) {
268 x_vis_[offsets_[t] + i] = x(i).vi_;
270 prepare_x_vis(args...);
273 template <
int R,
int C,
typename... Pargs>
275 const Pargs&... args) {
276 prepare_x_vis(args...);
279 template <
typename... Pargs>
281 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
282 for (
size_t i = 0; i < x.size(); ++i) {
283 x_vis_[offsets_[t] + i] = x[i].vi_;
285 prepare_x_vis(args...);
288 template <
typename... Pargs>
290 prepare_x_vis(args...);
293 template <
typename... Pargs>
295 prepare_x_vis(args...);
298 template <
typename... Pargs>
300 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
301 x_vis_[offsets_[t]] = x.
vi_;
302 prepare_x_vis(args...);
305 template <
typename... Pargs>
307 prepare_x_vis(args...);
310 template <
typename... Pargs>
312 prepare_x_vis(args...);
320 :
vari(
std::numeric_limits<double>::quiet_NaN()),
334 y_vi_[0] =
new vari(val_y,
false);
347 const std::vector<double>& val_y) {
348 M_[0] = val_y.size();
349 std::vector<var> var_y(M_[0]);
353 for (
size_t m = 0; m < var_y.size(); ++m) {
354 y_vi_[m] =
new vari(val_y[m],
false);
371 template <
int R,
int C>
373 const Eigen::Matrix<double, R, C>& val_y) {
374 M_[0] = val_y.rows();
375 M_[1] = val_y.cols();
376 Eigen::Matrix<var, R, C> var_y(M_[0], M_[1]);
380 for (
int m = 0; m < var_y.size(); ++m) {
381 y_vi_[m] =
new vari(val_y(m),
false);
409 count_memory(0, args...));
411 prepare_x_vis(args...);
413 return build_return_varis_and_vars(f_(is_var_,
value_of(args)...));
429 template <
int R,
int C,
typename... Pargs>
431 const Pargs&... args) {
432 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
434 for (
int n = 0; n < y_adj_jac.size(); ++n) {
435 x_vis_[offsets_[t] + n]->
adj_ += y_adj_jac(n);
439 accumulate_adjoints(args...);
453 template <
typename... Pargs>
455 const Pargs&... args) {
456 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
458 for (
size_t n = 0; n < y_adj_jac.size(); ++n)
459 x_vis_[offsets_[t] + n]->adj_ += y_adj_jac[n];
462 accumulate_adjoints(args...);
474 template <
typename... Pargs>
476 const Pargs&... args) {
477 accumulate_adjoints(args...);
491 template <
typename... Pargs>
493 static constexpr
int t =
sizeof...(Targs) -
sizeof...(Pargs) - 1;
495 x_vis_[offsets_[t]]->
adj_ += y_adj_jac;
498 accumulate_adjoints(args...);
510 template <
typename... Pargs>
512 accumulate_adjoints(args...);
532 auto y_adj_jacs = f_.multiply_adjoint_jacobian(is_var_, y_adj);
535 [
this](
auto&&... args) { this->accumulate_adjoints(args...); },
614 template <
typename F,
typename... Targs>
618 return (*vi)(args...);
size_t count_memory(size_t count, const int &x, const Pargs &... args)
size_t count_memory(size_t count, const std::vector< int > &x, const Pargs &... args)
void prepare_x_vis(const int &x, const Pargs &... args)
void prepare_x_vis(const double &x, const Pargs &... args)
void accumulate_adjoints()
size_t count_memory(size_t count, const Eigen::Matrix< double, R, C > &x, const Pargs &... args)
size_t count_memory(size_t count, const Eigen::Matrix< var, R, C > &x, const Pargs &... args)
count_memory returns count (the first argument) + the number of varis used in the second argument + t...
T value_of(const fvar< T > &v)
Return the value of the specified variable.
Defines a public enum named value which is defined to be false as the primitive scalar types cannot b...
std::array< int, internal::compute_dims< FReturnType >::value > M_
The variable implementation base class.
Compute the dimensionality of the given template argument.
void chain()
Propagate the adjoints at the output varis (y_vi_) back to the input varis (x_vis_) by: ...
var build_return_varis_and_vars(const double &val_y)
Return a var with a new vari holding the given value.
static STAN_THREADS_DEF AutodiffStackStorage * instance_
Independent (input) and dependent (output) variables for gradients.
auto operator()(const Targs &... args)
The adj_jac_vari functor.
(Expert) Numerical traits for algorithmic differentiation variables.
auto adj_jac_apply(const Targs &... args)
Return the result of applying the function defined by a nullary construction of F to the specified in...
adj_jac_vari interfaces a user supplied functor with the reverse mode autodiff.
void accumulate_adjoints(const double &y_adj_jac, const Pargs &... args)
Accumulate, if necessary, the value of y_adj_jac into the adjoint of the vari pointed to by the appro...
std::result_of_t< F(decltype(is_var_), decltype(value_of(Targs()))...)> FReturnType
void prepare_x_vis(const Eigen::Matrix< var, R, C > &x, const Pargs &... args)
prepare_x_vis populates x_vis_ with the varis from each of its input arguments.
void accumulate_adjoints(const Eigen::Matrix< double, R, C > &y_adj_jac, const Pargs &... args)
Accumulate, if necessary, the values of y_adj_jac into the adjoints of the varis pointed to by the ap...
void prepare_x_vis(const std::vector< int > &x, const Pargs &... args)
std::vector< var > build_return_varis_and_vars(const std::vector< double > &val_y)
Return a std::vector of vars created from newly allocated varis initialized with the values of val_y...
adj_jac_vari()
Initializes is_var_ with true if the scalar type in each argument is a var (and false if not) ...
size_t count_memory(size_t count, const std::vector< double > &x, const Pargs &... args)
size_t count_memory(size_t count)
vari * vi_
Pointer to the implementation of this variable.
void accumulate_adjoints(const std::vector< double > &y_adj_jac, const Pargs &... args)
Accumulate, if necessary, the values of y_adj_jac into the adjoints of the varis pointed to by the ap...
constexpr auto apply(const F &f, const Tuple &t)
Call the functor f with the tuple of arguments t, like:
constexpr auto apply_impl(const F &f, const Tuple &t, std::index_sequence< I... > i)
Invoke the functor f with arguments given in t and indexed in the index sequence I.
T * alloc_array(size_t n)
Allocate an array on the arena of the specified size to hold values of the specified template paramet...
size_t count_memory(size_t count, const var &x, const Pargs &... args)
void prepare_x_vis(const std::vector< double > &x, const Pargs &... args)
size_t count_memory(size_t count, const double &x, const Pargs &... args)
void prepare_x_vis(const var &x, const Pargs &... args)
double adj_
The adjoint of this variable, which is the partial derivative of this variable with respect to the ro...
void build_y_adj(vari **y_vi, const std::array< int, size > &M, double &y_adj)
Store the adjoint in y_vi[0] in y_adj.
size_t count_memory(size_t count, const std::vector< var > &x, const Pargs &... args)
void accumulate_adjoints(const int &y_adj_jac, const Pargs &... args)
Recursively call accumulate_adjoints with args.
void prepare_x_vis(const std::vector< var > &x, const Pargs &... args)
Eigen::Matrix< var, R, C > build_return_varis_and_vars(const Eigen::Matrix< double, R, C > &val_y)
Return an Eigen::Matrix of vars created from newly allocated varis initialized with the values of val...
void prepare_x_vis(const Eigen::Matrix< double, R, C > &x, const Pargs &... args)
void accumulate_adjoints(const std::vector< int > &y_adj_jac, const Pargs &... args)
Recursively call accumulate_adjoints with args.