vari.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_REV_VARI_HPP
#define STAN_MATH_OPENCL_REV_VARI_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/rev/core/vari.hpp>
#include <stan/math/prim/meta.hpp>
#include <stan/math/opencl/kernel_generator.hpp>
#include <stan/math/opencl/prim/reverse.hpp>
 
namespace stan {
namespace math {
 
template <typename T>
class vari_cl_base : public vari_base {
 public:
  T val_;
 
  T adj_;
 
  template <typename S, typename U,
            require_t<std::is_constructible<T, std::decay_t<S>>>* = nullptr,
            require_t<std::is_constructible<T, std::decay_t<U>>>* = nullptr>
  vari_cl_base(S&& val, U&& adj)
      : val_(std::forward<S>(val)), adj_(std::forward<U>(adj)) {}
 
  inline const auto& val() const { return val_; }
  inline auto& val_op() { return val_; }
 
  inline auto& adj() { return adj_; }
  inline auto& adj() const { return adj_; }
  inline auto& adj_op() { return adj_; }
 
  auto block(int row, int col, int rows, int cols) {
    auto&& val_block = stan::math::block_zero_based(val_, row, col, rows, cols);
    auto&& adj_block = stan::math::block_zero_based(adj_, row, col, rows, cols);
    return vari_view<std::decay_t<decltype(val_block)>>(std::move(val_block),
                                                        std::move(adj_block));
  }
 
  auto transpose() {
    auto&& val_t = stan::math::transpose(val_);
    auto&& adj_t = stan::math::transpose(adj_);
    return vari_view<std::decay_t<decltype(val_t)>>(std::move(val_t),
                                                    std::move(adj_t));
  }
 
  auto as_column_vector_or_scalar() {
    auto&& val_t = stan::math::as_column_vector_or_scalar(val_);
    auto&& adj_t = stan::math::as_column_vector_or_scalar(adj_);
    return vari_view<std::decay_t<decltype(val_t)>>(std::move(val_t),
                                                    std::move(adj_t));
  }
 
  auto reverse() {
    auto&& val_t = stan::math::reverse(val_);
    auto&& adj_t = stan::math::reverse(adj_);
    return vari_view<std::decay_t<decltype(val_t)>>(std::move(val_t),
                                                    std::move(adj_t));
  }
 
  template <typename RowIndex, typename ColIndex>
  auto index(const RowIndex& row_index, const ColIndex& col_index) {
    RowIndex r1 = row_index;
    RowIndex r2 = row_index;
    ColIndex c1 = col_index;
    ColIndex c2 = col_index;
    auto&& val_t = stan::math::indexing(val_, std::move(r1), std::move(c1));
    auto&& adj_t = stan::math::indexing(adj_, std::move(r2), std::move(c2));
    return vari_view<std::decay_t<decltype(val_t)>>(std::move(val_t),
                                                    std::move(adj_t));
  }
 
  const Eigen::Index rows() const { return val_.rows(); }
  const Eigen::Index cols() const { return val_.cols(); }
  const Eigen::Index size() const { return rows() * cols(); }
 
  virtual void chain() {}
};
 
template <typename T>
class vari_view<T, require_kernel_expression_lhs_t<T>>
    : public vari_cl_base<T> {
 public:
  static constexpr int RowsAtCompileTime{-1};
  static constexpr int ColsAtCompileTime{-1};
 
  using value_type = T;
  using vari_cl_base<T>::vari_cl_base;
  inline void set_zero_adjoint() final {}
};
 
template <typename T>
class vari_value<T, require_matrix_cl_t<T>> : public chainable_alloc,
                                              public vari_cl_base<T> {
 public:
  using value_type = T;
 
  static constexpr int RowsAtCompileTime{-1};
  static constexpr int ColsAtCompileTime{-1};
 
  template <typename S, require_convertible_t<S&, T>* = nullptr>
  explicit vari_value(S&& x)
      : chainable_alloc(),
        vari_cl_base<T>(std::forward<S>(x), constant(0, x.rows(), x.cols())) {
    ChainableStack::instance_->var_stack_.push_back(this);
  }
 
  template <typename S, require_eigen_t<S>* = nullptr,
            require_vt_same<T, S>* = nullptr>
  explicit vari_value(const S& x)
      : chainable_alloc(), vari_cl_base<T>(x, constant(0, x.rows(), x.cols())) {
    ChainableStack::instance_->var_nochain_stack_.push_back(this);
  }
 
  template <typename S, require_convertible_t<S&, T>* = nullptr>
  vari_value(S&& x, bool stacked)
      : chainable_alloc(),
        vari_cl_base<T>(std::forward<S>(x), constant(0, x.rows(), x.cols())) {
    if (stacked) {
      ChainableStack::instance_->var_stack_.push_back(this);
    } else {
      ChainableStack::instance_->var_nochain_stack_.push_back(this);
    }
  }
 
  template <typename S, typename K, require_convertible_t<T, S>* = nullptr,
            require_convertible_t<T, K>* = nullptr>
  explicit vari_value(S&& val, K&& adj)
      : chainable_alloc(),
        vari_cl_base<T>(std::forward<S>(val), std::forward<K>(adj)) {
    ChainableStack::instance_->var_nochain_stack_.push_back(this);
  }
 
  inline void set_zero_adjoint() final {
    this->adj_ = constant(0, this->rows(), this->cols());
  }
 
 private:
  template <typename, typename>
  friend class var_value;
};
 
}  // namespace math
}  // namespace stan
 
#endif
#endif