Stan Math Library  2.20.0
reverse mode automatic differentiation
adj_jac_apply.hpp
Go to the documentation of this file.
1 #ifndef STAN_MATH_REV_MAT_FUNCTOR_ADJ_JAC_APPLY_HPP
2 #define STAN_MATH_REV_MAT_FUNCTOR_ADJ_JAC_APPLY_HPP
3 
4 #include <stan/math/rev/meta.hpp>
8 #include <limits>
9 #include <tuple>
10 #include <vector>
11 
12 namespace stan {
13 namespace math {
14 
15 namespace internal {
28 template <class F, class Tuple, std::size_t... I>
29 constexpr auto apply_impl(const F& f, const Tuple& t,
30  std::index_sequence<I...> i) {
31  return f(std::get<I>(t)...);
32 }
33 
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>{}>{});
49 }
50 
59 template <size_t size>
60 void build_y_adj(vari** y_vi, const std::array<int, size>& M, double& y_adj) {
61  y_adj = y_vi[0]->adj_;
62 }
63 
72 template <size_t size>
73 void build_y_adj(vari** y_vi, const std::array<int, size>& M,
74  std::vector<double>& y_adj) {
75  y_adj.resize(M[0]);
76  for (size_t m = 0; m < y_adj.size(); ++m)
77  y_adj[m] = y_vi[m]->adj_;
78 }
79 
88 template <size_t size, int R, int C>
89 void build_y_adj(vari** y_vi, const std::array<int, size>& M,
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_;
94 }
95 
101 template <typename T>
102 struct compute_dims {};
103 
108 template <>
109 struct compute_dims<double> {
110  static constexpr size_t value = 0;
111 };
112 
117 template <typename T>
118 struct compute_dims<std::vector<T>> {
119  static constexpr size_t value = 1;
120 };
121 
126 template <typename T, int R, int C>
127 struct compute_dims<Eigen::Matrix<T, R, C>> {
128  static constexpr size_t value = 2;
129 };
130 } // namespace internal
131 
146 template <typename F, typename... Targs>
147 struct adj_jac_vari : public vari {
148  std::array<bool, sizeof...(Targs)> is_var_;
149  using FReturnType
150  = std::result_of_t<F(decltype(is_var_), decltype(value_of(Targs()))...)>;
151 
152  F f_;
153  std::array<int, sizeof...(Targs)> offsets_;
155  std::array<int, internal::compute_dims<FReturnType>::value> M_;
157 
181  template <int R, int C, typename... Pargs>
182  size_t count_memory(size_t count, const Eigen::Matrix<var, R, C>& x,
183  const Pargs&... args) {
184  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
185  offsets_[t] = count;
186  count += x.size();
187  return count_memory(count, args...);
188  }
189 
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;
194  offsets_[t] = count;
195  return count_memory(count, args...);
196  }
197 
198  template <typename... Pargs>
199  size_t count_memory(size_t count, const std::vector<var>& x,
200  const Pargs&... args) {
201  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
202  offsets_[t] = count;
203  count += x.size();
204  return count_memory(count, args...);
205  }
206 
207  template <typename... Pargs>
208  size_t count_memory(size_t count, const std::vector<double>& x,
209  const Pargs&... args) {
210  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
211  offsets_[t] = count;
212  return count_memory(count, args...);
213  }
214 
215  template <typename... Pargs>
216  size_t count_memory(size_t count, const std::vector<int>& x,
217  const Pargs&... args) {
218  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
219  offsets_[t] = count;
220  return count_memory(count, args...);
221  }
222 
223  template <typename... Pargs>
224  size_t count_memory(size_t count, const var& x, const Pargs&... args) {
225  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
226  offsets_[t] = count;
227  count += 1;
228  return count_memory(count, args...);
229  }
230 
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;
234  offsets_[t] = count;
235  return count_memory(count, args...);
236  }
237 
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;
241  offsets_[t] = count;
242  return count_memory(count, args...);
243  }
244 
245  size_t count_memory(size_t count) { return count; }
246 
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_;
269  }
270  prepare_x_vis(args...);
271  }
272 
273  template <int R, int C, typename... Pargs>
274  void prepare_x_vis(const Eigen::Matrix<double, R, C>& x,
275  const Pargs&... args) {
276  prepare_x_vis(args...);
277  }
278 
279  template <typename... Pargs>
280  void prepare_x_vis(const std::vector<var>& x, const Pargs&... args) {
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_;
284  }
285  prepare_x_vis(args...);
286  }
287 
288  template <typename... Pargs>
289  void prepare_x_vis(const std::vector<double>& x, const Pargs&... args) {
290  prepare_x_vis(args...);
291  }
292 
293  template <typename... Pargs>
294  void prepare_x_vis(const std::vector<int>& x, const Pargs&... args) {
295  prepare_x_vis(args...);
296  }
297 
298  template <typename... Pargs>
299  void prepare_x_vis(const var& x, const Pargs&... args) {
300  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
301  x_vis_[offsets_[t]] = x.vi_;
302  prepare_x_vis(args...);
303  }
304 
305  template <typename... Pargs>
306  void prepare_x_vis(const double& x, const Pargs&... args) {
307  prepare_x_vis(args...);
308  }
309 
310  template <typename... Pargs>
311  void prepare_x_vis(const int& x, const Pargs&... args) {
312  prepare_x_vis(args...);
313  }
314 
320  : vari(std::numeric_limits<double>::quiet_NaN()), // The val_ in this
321  // vari is unused
322  is_var_({{is_var<typename scalar_type<Targs>::type>::value...}}),
323  x_vis_(NULL),
324  y_vi_(NULL) {}
325 
332  var build_return_varis_and_vars(const double& val_y) {
334  y_vi_[0] = new vari(val_y, false);
335 
336  return y_vi_[0];
337  }
338 
346  std::vector<var> build_return_varis_and_vars(
347  const std::vector<double>& val_y) {
348  M_[0] = val_y.size();
349  std::vector<var> var_y(M_[0]);
350 
351  y_vi_
353  for (size_t m = 0; m < var_y.size(); ++m) {
354  y_vi_[m] = new vari(val_y[m], false);
355  var_y[m] = y_vi_[m];
356  }
357 
358  return var_y;
359  }
360 
371  template <int R, int C>
372  Eigen::Matrix<var, R, C> build_return_varis_and_vars(
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]);
377 
378  y_vi_
380  for (int m = 0; m < var_y.size(); ++m) {
381  y_vi_[m] = new vari(val_y(m), false);
382  var_y(m) = y_vi_[m];
383  }
384 
385  return var_y;
386  }
387 
388  void prepare_x_vis() {}
389 
407  auto operator()(const Targs&... args) {
409  count_memory(0, args...));
410 
411  prepare_x_vis(args...);
412 
413  return build_return_varis_and_vars(f_(is_var_, value_of(args)...));
414  }
415 
429  template <int R, int C, typename... Pargs>
430  void accumulate_adjoints(const Eigen::Matrix<double, R, C>& y_adj_jac,
431  const Pargs&... args) {
432  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
433  if (is_var_[t]) {
434  for (int n = 0; n < y_adj_jac.size(); ++n) {
435  x_vis_[offsets_[t] + n]->adj_ += y_adj_jac(n);
436  }
437  }
438 
439  accumulate_adjoints(args...);
440  }
441 
453  template <typename... Pargs>
454  void accumulate_adjoints(const std::vector<double>& y_adj_jac,
455  const Pargs&... args) {
456  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
457  if (is_var_[t]) {
458  for (size_t n = 0; n < y_adj_jac.size(); ++n)
459  x_vis_[offsets_[t] + n]->adj_ += y_adj_jac[n];
460  }
461 
462  accumulate_adjoints(args...);
463  }
464 
474  template <typename... Pargs>
475  void accumulate_adjoints(const std::vector<int>& y_adj_jac,
476  const Pargs&... args) {
477  accumulate_adjoints(args...);
478  }
479 
491  template <typename... Pargs>
492  void accumulate_adjoints(const double& y_adj_jac, const Pargs&... args) {
493  static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;
494  if (is_var_[t]) {
495  x_vis_[offsets_[t]]->adj_ += y_adj_jac;
496  }
497 
498  accumulate_adjoints(args...);
499  }
500 
510  template <typename... Pargs>
511  void accumulate_adjoints(const int& y_adj_jac, const Pargs&... args) {
512  accumulate_adjoints(args...);
513  }
514 
516 
528  void chain() {
529  FReturnType y_adj;
530 
531  internal::build_y_adj(y_vi_, M_, y_adj);
532  auto y_adj_jacs = f_.multiply_adjoint_jacobian(is_var_, y_adj);
533 
535  [this](auto&&... args) { this->accumulate_adjoints(args...); },
536  y_adj_jacs);
537  }
538 };
539 
614 template <typename F, typename... Targs>
615 auto adj_jac_apply(const Targs&... args) {
616  auto vi = new adj_jac_vari<F, Targs...>();
617 
618  return (*vi)(args...);
619 }
620 
621 } // namespace math
622 } // namespace stan
623 #endif
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)
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.
Definition: value_of.hpp:17
Defines a public enum named value which is defined to be false as the primitive scalar types cannot b...
Definition: is_var.hpp:10
std::array< int, internal::compute_dims< FReturnType >::value > M_
The variable implementation base class.
Definition: vari.hpp:30
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.
Definition: var.hpp:33
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.
Definition: var.hpp:45
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...
Definition: vari.hpp:44
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.

     [ Stan Home Page ] © 2011–2018, Stan Development Team.