ordered_logistic_lpmf.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_PRIM_ORDERED_LOGISTIC_LPMF_HPP
#define STAN_MATH_OPENCL_PRIM_ORDERED_LOGISTIC_LPMF_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/kernel_generator.hpp>
#include <stan/math/opencl/kernels/add.hpp>
#include <stan/math/opencl/kernels/ordered_logistic_lpmf.hpp>
#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/err.hpp>
#include <stan/math/prim/fun/constants.hpp>
#include <stan/math/prim/fun/elt_divide.hpp>
#include <stan/math/prim/fun/elt_multiply.hpp>
#include <stan/math/prim/functor/partials_propagator.hpp>
#include <stan/math/prim/err/constraint_tolerance.hpp>
 
namespace stan {
namespace math {
 
template <bool propto, typename T_y_cl, typename T_loc_cl, typename T_cuts_cl,
          require_all_prim_or_rev_kernel_expression_t<T_y_cl, T_loc_cl,
                                                      T_cuts_cl>* = nullptr>
inline return_type_t<T_y_cl, T_loc_cl, T_cuts_cl> ordered_logistic_lpmf(
    const T_y_cl& y, const T_loc_cl& lambda, const T_cuts_cl& cuts) {
  constexpr bool is_y_vector = !is_stan_scalar<T_y_cl>::value;
  static constexpr const char* function = "ordered_logistic_lpmf(OpenCL)";
 
  if (size(y) != 1) {
    check_size_match(function, "Size of ", "y", math::size(y), "Size of",
                     "lambda", math::size(lambda));
  }
 
  int N_instances = max_size(y, lambda);
  int N_classes = cuts.rows() + 1;
  int N_cut_sets = cuts.cols();
 
  if (N_cut_sets > 1) {
    check_size_match(function, "Length of lambda variables ", N_instances,
                     "Number of cutpoint vectors ", N_cut_sets);
  }
  if (N_instances == 0 || N_classes == 1) {
    return 0.0;
  }
  const auto& cuts_val = eval(value_of(cuts));
  if (N_classes >= 2) {
    auto cuts_head
        = block_zero_based(cuts_val, 0, 0, cuts.rows() - 1, N_cut_sets);
    auto cuts_tail
        = block_zero_based(cuts_val, 1, 0, cuts.rows() - 1, N_cut_sets);
    check_cl(function, "Cuts", cuts_head, "ordered and finite")
        = cuts_head < cuts_tail && isfinite(cuts_head) && isfinite(cuts_tail);
  } else if (N_classes == 1) {
    check_cl(function, "Cuts", cuts_val, "finite") = isfinite(cuts_val);
  }
 
  if constexpr (!include_summand<propto, T_loc_cl, T_cuts_cl>::value) {
    return 0.0;
  }
 
  const auto& y_val = eval(value_of(y));
  const auto& lambda_val = eval(value_of(lambda));
 
  const auto& y_val_cl = to_matrix_cl(y_val);
 
  const int local_size
      = opencl_kernels::ordered_logistic.get_option("LOCAL_SIZE_");
  const int wgs = (N_instances + local_size - 1) / local_size;
 
  bool need_lambda_derivative = is_autodiff_v<T_loc_cl>;
  bool need_cuts_derivative = is_autodiff_v<T_cuts_cl>;
  bool need_broadcasting = N_cut_sets == 1 && N_instances != 1;
  matrix_cl<double> logp_cl(wgs, 1);
  matrix_cl<double> lambda_derivative_cl(N_instances,
                                         need_lambda_derivative ? 1 : 0);
  matrix_cl<double> cuts_derivative_cl(
      N_classes - 1,
      need_cuts_derivative ? (need_broadcasting ? wgs : N_cut_sets) : 0);
 
  try {
    opencl_kernels::ordered_logistic(
        cl::NDRange(local_size * wgs), cl::NDRange(local_size), logp_cl,
        lambda_derivative_cl, cuts_derivative_cl, y_val_cl, lambda_val,
        cuts_val, N_instances, N_classes, is_y_vector, !need_broadcasting,
        need_lambda_derivative, need_cuts_derivative);
  } catch (const cl::Error& e) {
    check_opencl_error(function, e);
  }
 
  double logp = sum(from_matrix_cl(logp_cl));
 
  if (!std::isfinite(logp)) {
    results(check_cl(function, "Vector of dependent variables", y_val,
                     "between 0 and number of classes"),
            check_cl(function, "lambda vector", lambda_val, "finite"))
        = expressions(y_val >= 1 && y_val <= static_cast<int>(N_classes),
                      isfinite(lambda_val));
  }
  auto ops_partials = make_partials_propagator(lambda, cuts);
 
  if constexpr (is_autodiff_v<T_loc_cl>) {
    partials<0>(ops_partials) = lambda_derivative_cl;
  }
  if constexpr (is_autodiff_v<T_cuts_cl>) {
    if (need_broadcasting) {
      partials<1>(ops_partials) = rowwise_sum(cuts_derivative_cl);
    } else {
      partials<1>(ops_partials) = std::move(cuts_derivative_cl);
    }
  }
  return ops_partials.build(logp);
}
 
}  // namespace math
}  // namespace stan
#endif
#endif