rep_matrix.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_REV_REP_MATRIX_HPP
#define STAN_MATH_OPENCL_REV_REP_MATRIX_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/kernel_generator.hpp>
#include <stan/math/rev/core.hpp>
#include <stan/math/rev/fun/value_of.hpp>
#include <stan/math/opencl/kernels/rep_matrix.hpp>
#include <stan/math/opencl/prim/rep_matrix.hpp>
#include <stan/math/opencl/prim/sum.hpp>
 
namespace stan {
namespace math {
 
template <typename T_ret, require_var_vt<is_matrix_cl, T_ret>* = nullptr>
inline var_value<matrix_cl<double>> rep_matrix(const var& A, int n, int m) {
  return make_callback_var(rep_matrix<matrix_cl<double>>(A.val(), n, m),
                           [A](vari_value<matrix_cl<double>>& res) mutable {
                             A.adj() += sum(res.adj());
                           });
}
 
template <typename T,
          require_all_kernel_expressions_and_none_scalar_t<T>* = nullptr>
inline var_value<matrix_cl<double>> rep_matrix(const var_value<T>& A, int m) {
  return make_callback_var(
      rep_matrix(A.val(), m), [A](vari_value<matrix_cl<double>>& res) mutable {
        if (A.adj().size() != 0) {
          matrix_cl<double> A_adj = std::move(A.adj());
          try {
            opencl_kernels::rep_matrix_rev(
                cl::NDRange(A_adj.rows(), A_adj.cols()), A_adj, res.adj(),
                res.adj().rows(), res.adj().cols(), res.adj().view());
          } catch (const cl::Error& e) {
            check_opencl_error("rep_matrix(rev OpenCL)", e);
          }
          A.adj() = std::move(A_adj);
        }
      });
}
 
}  // namespace math
}  // namespace stan
 
#endif
#endif