var.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_REV_CORE_VAR_HPP
#define STAN_MATH_REV_CORE_VAR_HPP
 
#ifdef STAN_OPENCL
#include <stan/math/opencl/rev/vari.hpp>
#include <stan/math/opencl/plain_type.hpp>
#endif
#include <stan/math/rev/core/vari.hpp>
#include <stan/math/rev/core/grad.hpp>
#include <stan/math/rev/core/chainable_alloc.hpp>
#include <stan/math/prim/meta.hpp>
#include <stan/math/rev/meta/is_vari.hpp>
#include <stan/math/rev/meta/arena_type.hpp>
#include <stan/math/rev/core/reverse_pass_callback.hpp>
#include <ostream>
#include <vector>
 
namespace stan {
namespace math {
 
// forward declare
template <typename Vari>
static void grad(Vari* vi);
 
template <typename T>
class var_value<T, require_floating_point_t<T>> {
 public:
  using value_type = std::decay_t<T>;  // type in vari_value.
  using vari_type = vari_value<value_type>;
 
  vari_type* vi_;
 
  inline bool is_uninitialized() { return (vi_ == nullptr); }
 
  var_value() : vi_(nullptr) {}
 
  template <typename S, require_convertible_t<S&, value_type>* = nullptr>
  var_value(S x) : vi_(new vari_type(x, false)) {}  // NOLINT
 
  var_value(vari_type* vi) : vi_(vi) {}  // NOLINT
 
  inline const auto& val() const noexcept { return vi_->val(); }
 
  inline auto& adj() const noexcept { return vi_->adj(); }
 
  inline auto& adj() noexcept { return vi_->adj_; }
 
  inline void grad(std::vector<var_value<T>>& x, std::vector<value_type>& g) {
    stan::math::grad(vi_);
    g.resize(x.size());
    for (size_t i = 0; i < x.size(); ++i) {
      g[i] = x[i].vi_->adj_;
    }
  }
 
  void grad() { stan::math::grad(vi_); }
 
  // POINTER OVERRIDES
 
  inline vari_type& operator*() { return *vi_; }
 
  inline vari_type* operator->() { return vi_; }
 
  // COMPOUND ASSIGNMENT OPERATORS
 
  inline var_value<T>& operator+=(const var_value<T>& b);
 
  inline var_value<T>& operator+=(T b);
 
  inline var_value<T>& operator-=(const var_value<T>& b);
 
  inline var_value<T>& operator-=(T b);
 
  inline var_value<T>& operator*=(const var_value<T>& b);
 
  inline var_value<T>& operator*=(T b);
 
  inline var_value<T>& operator/=(const var_value<T>& b);
 
  inline var_value<T>& operator/=(T b);
 
  friend std::ostream& operator<<(std::ostream& os, const var_value<T>& v) {
    if (v.vi_ == nullptr) {
      return os << "uninitialized";
    }
    return os << v.val();
  }
};
 
namespace internal {
template <typename T>
using require_matrix_var_value = require_t<bool_constant<
    (is_eigen<T>::value || is_kernel_expression_and_not_scalar<T>::value)
    && std::is_floating_point<value_type_t<T>>::value>>;
}
 
template <typename T>
class var_value<T, internal::require_matrix_var_value<T>> {
 public:
  using value_type = T;  // type in vari_value.
  using vari_type = std::conditional_t<is_plain_type<value_type>::value,
                                       vari_value<value_type>, vari_view<T>>;
 
  static constexpr int RowsAtCompileTime{vari_type::RowsAtCompileTime};
  static constexpr int ColsAtCompileTime{vari_type::ColsAtCompileTime};
 
  vari_type* vi_;
 
  inline bool is_uninitialized() noexcept { return (vi_ == nullptr); }
 
  var_value() : vi_(nullptr) {}
 
  template <typename S, require_assignable_t<value_type, S>* = nullptr>
  var_value(S&& x) : vi_(new vari_type(std::forward<S>(x), false)) {}  // NOLINT
 
  template <typename S,
            require_assignable_t<vari_type,
                                 typename var_value<S>::vari_type>* = nullptr,
            require_all_plain_type_t<T, S>* = nullptr>
  var_value(const var_value<S>& other) : vi_(other.vi_) {}
 
  template <typename S,
            require_not_assignable_t<
                vari_type, typename var_value<S>::vari_type>* = nullptr,
            require_constructible_t<vari_type, S>* = nullptr,
            require_all_plain_type_t<T, S>* = nullptr>
  var_value(const var_value<S>& other) : vi_(new vari_type(other.vi_->val_)) {
    reverse_pass_callback([this_vi = this->vi_, other_vi = other.vi_]() {
      other_vi->adj_ += this_vi->adj_;
    });
  }
 
  template <typename S, typename T_ = T,
            require_assignable_t<value_type, S>* = nullptr,
            require_not_plain_type_t<S>* = nullptr,
            require_plain_type_t<T_>* = nullptr>
  var_value(const var_value<S>& other) : vi_(new vari_type(other.vi_->val_)) {
    reverse_pass_callback(
        [this_vi = this->vi_, other_vi = other.vi_]() mutable {
          other_vi->adj_ += this_vi->adj_;
        });
  }
 
  template <typename S, typename T_ = T,
            require_assignable_t<value_type, S>* = nullptr,
            require_arena_matrix_t<S>* = nullptr>
  var_value(const S& val, const S& adj) : vi_(new vari_type(val, adj)) {}
 
  var_value(vari_type* vi) : vi_(vi) {}  // NOLINT
 
  inline const auto& val() const noexcept { return vi_->val(); }
  inline auto& val_op() noexcept { return vi_->val_op(); }
 
  inline auto& adj() noexcept { return vi_->adj(); }
  inline auto& adj() const noexcept { return vi_->adj(); }
  inline auto& adj_op() noexcept { return vi_->adj(); }
 
  inline Eigen::Index rows() const noexcept { return vi_->rows(); }
  inline Eigen::Index cols() const noexcept { return vi_->cols(); }
  inline Eigen::Index size() const noexcept { return vi_->size(); }
 
  // POINTER OVERRIDES
 
  inline vari_type& operator*() { return *vi_; }
 
  inline vari_type* operator->() { return vi_; }
 
  // COMPOUND ASSIGNMENT OPERATORS
 
  inline var_value<T>& operator+=(const var_value<T>& b);
 
  inline var_value<T>& operator+=(T b);
 
  template <typename S, require_st_var<S>* = nullptr>
  inline var_value<T>& operator-=(const S& b);
 
  template <typename S, require_st_arithmetic<S>* = nullptr>
  inline var_value<T>& operator-=(const S& b);
 
  inline var_value<T>& operator*=(const var_value<T>& b);
 
  inline var_value<T>& operator*=(T b);
 
  inline var_value<T>& operator/=(const var_value<T>& b);
 
  inline var_value<T>& operator/=(T b);
 
  inline auto block(Eigen::Index start_row, Eigen::Index start_col,
                    Eigen::Index num_rows, Eigen::Index num_cols) const {
    using vari_sub
        = decltype(vi_->block(start_row, start_col, num_rows, num_cols));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(
        new vari_sub(vi_->block(start_row, start_col, num_rows, num_cols)));
  }
  inline auto block(Eigen::Index start_row, Eigen::Index start_col,
                    Eigen::Index num_rows, Eigen::Index num_cols) {
    using vari_sub
        = decltype(vi_->block(start_row, start_col, num_rows, num_cols));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(
        new vari_sub(vi_->block(start_row, start_col, num_rows, num_cols)));
  }
 
  inline auto transpose() const {
    using vari_sub = decltype(vi_->transpose());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->transpose()));
  }
  inline auto transpose() {
    using vari_sub = decltype(vi_->transpose());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->transpose()));
  }
 
  inline auto head(Eigen::Index n) const {
    using vari_sub = decltype(vi_->head(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->head(n)));
  }
  inline auto head(Eigen::Index n) {
    using vari_sub = decltype(vi_->head(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->head(n)));
  }
 
  inline auto tail(Eigen::Index n) const {
    using vari_sub = decltype(vi_->tail(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->tail(n)));
  }
  inline auto tail(Eigen::Index n) {
    using vari_sub = decltype(vi_->tail(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->tail(n)));
  }
 
  inline auto segment(Eigen::Index i, Eigen::Index n) const {
    using vari_sub = decltype(vi_->segment(i, n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->segment(i, n)));
  }
  inline auto segment(Eigen::Index i, Eigen::Index n) {
    using vari_sub = decltype(vi_->segment(i, n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->segment(i, n)));
  }
 
  inline auto row(Eigen::Index i) const {
    using vari_sub = decltype(vi_->row(i));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->row(i)));
  }
  inline auto row(Eigen::Index i) {
    using vari_sub = decltype(vi_->row(i));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->row(i)));
  }
 
  inline auto col(Eigen::Index i) const {
    using vari_sub = decltype(vi_->col(i));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->col(i)));
  }
  inline auto col(Eigen::Index i) {
    using vari_sub = decltype(vi_->col(i));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->col(i)));
  }
 
  inline auto diagonal() const {
    using vari_sub = decltype(vi_->diagonal());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->diagonal()));
  }
  inline auto diagonal() {
    using vari_sub = decltype(vi_->diagonal());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->diagonal()));
  }
 
  inline auto as_column_vector_or_scalar() const {
    using vari_sub = decltype(vi_->as_column_vector_or_scalar());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->as_column_vector_or_scalar()));
  }
  inline auto as_column_vector_or_scalar() {
    using vari_sub = decltype(vi_->as_column_vector_or_scalar());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->as_column_vector_or_scalar()));
  }
 
  inline auto coeff(Eigen::Index i) const {
    using vari_sub = decltype(vi_->coeff(i));
    vari_sub* vari_coeff = new vari_sub(vi_->coeff(i));
    reverse_pass_callback([this_vi = this->vi_, vari_coeff, i]() {
      this_vi->adj_.coeffRef(i) += vari_coeff->adj_;
    });
    return var_value<value_type_t<vari_sub>>(vari_coeff);
  }
  inline auto coeff(Eigen::Index i) {
    using vari_sub = decltype(vi_->coeff(i));
    vari_sub* vari_coeff = new vari_sub(vi_->coeff(i));
    reverse_pass_callback([this_vi = this->vi_, vari_coeff, i]() {
      this_vi->adj_.coeffRef(i) += vari_coeff->adj_;
    });
    return var_value<value_type_t<vari_sub>>(vari_coeff);
  }
 
  inline auto coeff(Eigen::Index i, Eigen::Index j) const {
    using vari_sub = decltype(vi_->coeff(i, j));
    vari_sub* vari_coeff = new vari_sub(vi_->coeff(i, j));
    reverse_pass_callback([this_vi = this->vi_, vari_coeff, i, j]() {
      this_vi->adj_.coeffRef(i, j) += vari_coeff->adj_;
    });
    return var_value<value_type_t<vari_sub>>(vari_coeff);
  }
  inline auto coeff(Eigen::Index i, Eigen::Index j) {
    using vari_sub = decltype(vi_->coeff(i, j));
    vari_sub* vari_coeff = new vari_sub(vi_->coeff(i, j));
    reverse_pass_callback([this_vi = this->vi_, vari_coeff, i, j]() {
      this_vi->adj_.coeffRef(i, j) += vari_coeff->adj_;
    });
    return var_value<value_type_t<vari_sub>>(vari_coeff);
  }
 
  inline auto operator()(Eigen::Index i) const { return this->coeff(i); }
  inline auto operator()(Eigen::Index i) { return this->coeff(i); }
 
  inline auto operator()(Eigen::Index i, Eigen::Index j) const {
    return this->coeff(i, j);
  }
  inline auto operator()(Eigen::Index i, Eigen::Index j) {
    return this->coeff(i, j);
  }
 
  inline auto coeffRef(Eigen::Index i) const { return this->coeff(i); }
  inline auto coeffRef(Eigen::Index i) { return this->coeff(i); }
 
  inline auto coeffRef(Eigen::Index i, Eigen::Index j) const {
    return this->coeff(i, j);
  }
  inline auto coeffRef(Eigen::Index i, Eigen::Index j) {
    return this->coeff(i, j);
  }
 
  inline auto rowwise_reverse() const {
    using vari_sub = decltype(vi_->rowwise_reverse());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->rowwise_reverse()));
  }
  inline auto rowwise_reverse() {
    using vari_sub = decltype(vi_->rowwise_reverse());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->rowwise_reverse()));
  }
 
  inline auto colwise_reverse() const {
    using vari_sub = decltype(vi_->colwise_reverse());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->colwise_reverse()));
  }
  inline auto colwise_reverse() {
    using vari_sub = decltype(vi_->colwise_reverse());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->colwise_reverse()));
  }
 
  inline auto reverse() const {
    using vari_sub = decltype(vi_->reverse());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->reverse()));
  }
  inline auto reverse() {
    using vari_sub = decltype(vi_->reverse());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->reverse()));
  }
 
  inline auto topRows(Eigen::Index n) const {
    using vari_sub = decltype(vi_->topRows(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->topRows(n)));
  }
  inline auto topRows(Eigen::Index n) {
    using vari_sub = decltype(vi_->topRows(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->topRows(n)));
  }
 
  inline auto bottomRows(Eigen::Index n) const {
    using vari_sub = decltype(vi_->bottomRows(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->bottomRows(n)));
  }
  inline auto bottomRows(Eigen::Index n) {
    using vari_sub = decltype(vi_->bottomRows(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->bottomRows(n)));
  }
 
  inline auto middleRows(Eigen::Index start_row, Eigen::Index n) const {
    using vari_sub = decltype(vi_->middleRows(start_row, n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->middleRows(start_row, n)));
  }
  inline auto middleRows(Eigen::Index start_row, Eigen::Index n) {
    using vari_sub = decltype(vi_->middleRows(start_row, n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->middleRows(start_row, n)));
  }
 
  inline auto leftCols(Eigen::Index n) const {
    using vari_sub = decltype(vi_->leftCols(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->leftCols(n)));
  }
  inline auto leftCols(Eigen::Index n) {
    using vari_sub = decltype(vi_->leftCols(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->leftCols(n)));
  }
 
  inline auto rightCols(Eigen::Index n) const {
    using vari_sub = decltype(vi_->rightCols(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->rightCols(n)));
  }
  inline auto rightCols(Eigen::Index n) {
    using vari_sub = decltype(vi_->rightCols(n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->rightCols(n)));
  }
 
  inline auto middleCols(Eigen::Index start_col, Eigen::Index n) const {
    using vari_sub = decltype(vi_->middleCols(start_col, n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->middleCols(start_col, n)));
  }
  inline auto middleCols(Eigen::Index start_col, Eigen::Index n) {
    using vari_sub = decltype(vi_->middleCols(start_col, n));
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->middleCols(start_col, n)));
  }
 
  inline auto array() const {
    using vari_sub = decltype(vi_->array());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->array()));
  }
  inline auto array() {
    using vari_sub = decltype(vi_->array());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->array()));
  }
 
  inline auto matrix() const {
    using vari_sub = decltype(vi_->matrix());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->matrix()));
  }
  inline auto matrix() {
    using vari_sub = decltype(vi_->matrix());
    using var_sub = var_value<value_type_t<vari_sub>>;
    return var_sub(new vari_sub(vi_->matrix()));
  }
 
  friend std::ostream& operator<<(std::ostream& os, const var_value<T>& v) {
    if (v.vi_ == nullptr) {
      return os << "uninitialized";
    }
    return os << v.val();
  }
 
  template <typename U = T,
            require_any_t<is_eigen<U>, is_matrix_cl<U>>* = nullptr>
  inline auto rows() const noexcept {
    return vi_->rows();
  }
 
  template <typename U = T,
            require_any_t<is_eigen<U>, is_matrix_cl<U>>* = nullptr>
  inline auto cols() const noexcept {
    return vi_->cols();
  }
 
  template <typename S, require_assignable_t<value_type, S>* = nullptr,
            require_all_plain_type_t<T, S>* = nullptr,
            require_same_t<plain_type_t<T>, plain_type_t<S>>* = nullptr>
  inline var_value<T>& operator=(const var_value<S>& other) {
    vi_ = other.vi_;
    return *this;
  }
 
  template <typename S, typename T_ = T,
            require_assignable_t<value_type, S>* = nullptr,
            require_all_plain_type_t<T_, S>* = nullptr,
            require_not_same_t<plain_type_t<T_>, plain_type_t<S>>* = nullptr>
  inline var_value<T>& operator=(const var_value<S>& other) {
    static_assert(
        EIGEN_PREDICATE_SAME_MATRIX_SIZE(T, S),
        "You mixed matrices of different sizes that are not assignable.");
    vi_ = new vari_type(other.vi_);
    return *this;
  }
 
  template <typename S, typename T_ = T,
            require_assignable_t<value_type, S>* = nullptr,
            require_not_plain_type_t<S>* = nullptr,
            require_plain_type_t<T_>* = nullptr>
  inline var_value<T>& operator=(const var_value<S>& other) {
    // If vi_ is nullptr then the var needs initialized via copy constructor
    if (!(this->vi_)) {
      *this = var_value<T>(other);
      return *this;
    }
    arena_t<plain_type_t<T>> prev_val(vi_->val_.rows(), vi_->val_.cols());
    prev_val.deep_copy(vi_->val_);
    vi_->val_.deep_copy(other.val());
    // no need to change any adjoints - these are just zeros before the reverse
    // pass
 
    reverse_pass_callback(
        [this_vi = this->vi_, other_vi = other.vi_, prev_val]() mutable {
          this_vi->val_.deep_copy(prev_val);
 
          // we have no way of detecting aliasing between this->vi_->adj_ and
          // other.vi_->adj_, so we must copy adjoint before reseting to zero
 
          // we can reuse prev_val instead of allocating a new matrix
          prev_val.deep_copy(this_vi->adj_);
          this_vi->adj_.setZero();
          other_vi->adj_ += prev_val;
        });
    return *this;
  }
  template <typename S, typename T_ = T,
            require_assignable_t<value_type, S>* = nullptr,
            require_not_plain_type_t<T_>* = nullptr>
  inline var_value<T>& operator=(const var_value<S>& other) {
    // If vi_ is nullptr then the var needs initialized via copy constructor
    if (!(this->vi_)) {
      []() STAN_COLD_PATH {
        throw std::domain_error(
            "var_value<matrix>::operator=(var_value<expression>):"
            " Internal Bug! Please report this with an example"
            " of your model to the Stan math github repository.");
      }();
    }
    arena_t<plain_type_t<T>> prev_val = vi_->val_;
    vi_->val_ = other.val();
    // no need to change any adjoints - these are just zeros before the reverse
    // pass
 
    reverse_pass_callback(
        [this_vi = this->vi_, other_vi = other.vi_, prev_val]() mutable {
          this_vi->val_ = prev_val;
 
          // we have no way of detecting aliasing between this->vi_->adj_ and
          // other.vi_->adj_, so we must copy adjoint before reseting to zero
 
          // we can reuse prev_val instead of allocating a new matrix
          prev_val = this_vi->adj_;
          this_vi->adj_.setZero();
          other_vi->adj_ += prev_val;
        });
    return *this;
  }
 
  template <typename T_ = T, require_plain_type_t<T_>* = nullptr>
  inline auto& eval() noexcept {
    return *this;
  }
  template <typename T_ = T, require_plain_type_t<T_>* = nullptr>
  inline const auto& eval() const noexcept {
    return *this;
  }
 
  template <typename T_ = T, require_not_plain_type_t<T_>* = nullptr>
  inline auto eval() {
    return var_value<plain_type_t<T>>(*this);
  }
  template <typename T_ = T, require_not_plain_type_t<T_>* = nullptr>
  inline auto eval() const {
    return var_value<plain_type_t<T>>(*this);
  }
 
  inline var_value<T>& operator=(const var_value<T>& other) {
    return operator=<T>(other);
  }
};
 
// For backwards compatability the default value is double
using var = var_value<double>;
 
}  // namespace math
 
template <typename T>
struct scalar_type<T, std::enable_if_t<is_var<T>::value>> {
  using type
      = math::var_value<scalar_type_t<typename std::decay_t<T>::value_type>>;
};
 
}  // namespace stan
#endif