precomputed_gradients.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CORE_PRECOMPUTED_GRADIENTS_HPP
#define STAN_MATH_REV_CORE_PRECOMPUTED_GRADIENTS_HPP
 
#include <stan/math/prim/err/check_matching_dims.hpp>
#include <stan/math/rev/core/vari.hpp>
#include <stan/math/rev/core/var.hpp>
#include <stan/math/rev/fun/dims.hpp>
#include <stan/math/rev/fun/to_arena.hpp>
#include <algorithm>
#include <vector>
#include <tuple>
 
namespace stan {
namespace math {
 
template <typename ContainerOperands = std::tuple<>,
          typename ContainerGradients = std::tuple<>>
class precomputed_gradients_vari_template : public vari {
 protected:
  const size_t size_;
  vari** varis_;
  double* gradients_;
  static_assert(std::tuple_size<ContainerOperands>::value
                    == std::tuple_size<ContainerGradients>::value,
                "precomputed_gradients_vari: ContainerOperands and "
                "ContainerGradients should have same size!");
  static constexpr size_t N_containers
      = std::tuple_size<ContainerOperands>::value;
  ContainerOperands container_operands_;
  ContainerGradients container_gradients_;
 
  template <size_t... Is>
  void check_sizes(std::index_sequence<Is...>) {
    static_cast<void>(std::initializer_list<int>{
        (check_matching_dims("precomputed_gradients_vari", "operands",
                             std::get<Is>(container_operands_), "gradients",
                             std::get<Is>(container_gradients_)),
         0)...});
  }
 
 public:
  template <typename... ContainerOps, typename... ContainerGrads>
  precomputed_gradients_vari_template(
      double val, size_t size, vari** varis, double* gradients,
      const std::tuple<ContainerOps...>& container_operands = std::tuple<>(),
      const std::tuple<ContainerGrads...>& container_gradients = std::tuple<>())
      : vari(val),
        size_(size),
        varis_(varis),
        gradients_(gradients),
        container_operands_(index_apply<N_containers>([&](auto... Is) {
          return std::make_tuple(to_arena(std::get<Is>(container_operands))...);
        })),
        container_gradients_(index_apply<N_containers>([&](auto... Is) {
          return std::make_tuple(
              to_arena(std::get<Is>(container_gradients))...);
        })) {
    check_sizes(std::make_index_sequence<N_containers>());
  }
 
  template <typename Arith, typename VecVar, typename VecArith,
            typename... ContainerOps, typename... ContainerGrads,
            require_all_vector_t<VecVar, VecArith>* = nullptr>
  precomputed_gradients_vari_template(
      Arith val, const VecVar& vars, const VecArith& gradients,
      const std::tuple<ContainerOps...>& container_operands = std::tuple<>(),
      const std::tuple<ContainerGrads...>& container_gradients = std::tuple<>())
      : vari(val),
        size_(vars.size()),
        varis_(ChainableStack::instance_->memalloc_.alloc_array<vari*>(
            vars.size())),
        gradients_(ChainableStack::instance_->memalloc_.alloc_array<double>(
            vars.size())),
        container_operands_(index_apply<N_containers>([&](auto... Is) {
          return std::make_tuple(to_arena(std::get<Is>(container_operands))...);
        })),
        container_gradients_(index_apply<N_containers>([&](auto... Is) {
          return std::make_tuple(
              to_arena(std::get<Is>(container_gradients))...);
        })) {
    check_consistent_sizes("precomputed_gradients_vari", "vars", vars,
                           "gradients", gradients);
    check_sizes(std::make_index_sequence<N_containers>());
    for (size_t i = 0; i < vars.size(); ++i) {
      varis_[i] = vars[i].vi_;
    }
    std::copy(gradients.begin(), gradients.end(), gradients_);
  }
 
  void chain() {
    for (size_t i = 0; i < size_; ++i) {
      varis_[i]->adj_ += adj_ * gradients_[i];
    }
    index_apply<N_containers>([this](auto... Is) {
      static_cast<void>(std::initializer_list<int>{
          (this->chain_one(std::get<Is>(this->container_operands_),
                           std::get<Is>(this->container_gradients_)),
           0)...});
    });
  }
 
 private:
  template <typename Op, typename Grad,
            require_all_not_std_vector_t<Op, Grad>* = nullptr>
  void chain_one(Op& op, const Grad& grad) {
    op.vi_->adj_ = op.vi_->adj_ + this->adj_ * grad;
  }
 
  template <typename Op, typename OpAlloc, typename Grad, typename GradAlloc>
  void chain_one(const std::vector<Op, OpAlloc>& op,
                 const std::vector<Grad, GradAlloc>& grad) {
    for (int i = 0; i < op.size(); i++) {
      chain_one(op[i], grad[i]);
    }
  }
};
 
using precomputed_gradients_vari
    = precomputed_gradients_vari_template<std::tuple<>, std::tuple<>>;
 
template <typename Arith, typename VecVar, typename VecArith,
          typename... ContainerOperands, typename... ContainerGradients>
inline var precomputed_gradients(
    Arith value, const VecVar& operands, const VecArith& gradients,
    const std::tuple<ContainerOperands...>& container_operands = std::tuple<>(),
    const std::tuple<ContainerGradients...>& container_gradients
    = std::tuple<>()) {
  return {new precomputed_gradients_vari_template<
      std::tuple<arena_t<ContainerOperands>...>,
      std::tuple<arena_t<ContainerGradients>...>>(
      value, operands, gradients, container_operands, container_gradients)};
}
 
}  // namespace math
}  // namespace stan
#endif