log_mix.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_OPENCL_PRIM_LOG_MIX_HPP
#define STAN_MATH_OPENCL_PRIM_LOG_MIX_HPP
#ifdef STAN_OPENCL
 
#include <stan/math/opencl/matrix_cl.hpp>
#include <stan/math/opencl/kernel_generator.hpp>
#include <stan/math/opencl/rev/operands_and_partials.hpp>
 
namespace stan {
namespace math {
 
template <typename T_theta_cl, typename T_lambda_cl,
          require_all_prim_or_rev_kernel_expression_t<T_theta_cl,
                                                      T_lambda_cl>* = nullptr>
inline auto log_mix(const T_theta_cl& theta, const T_lambda_cl& lambda) {
  static constexpr const char* function = "log_mix(OpenCL)";
  using T_return = return_type_t<T_theta_cl, T_lambda_cl>;
  using std::isfinite;
  using std::isnan;
 
  check_consistent_sizes(function, "theta", theta, "lambda", lambda);
  const size_t N = max_size(theta, lambda);
  if (N == 0) {
    return T_return(0.0);
  }
 
  const auto& theta_col = as_column_vector_or_scalar(theta);
 
  const auto& theta_val = value_of(theta_col);
  const auto& lambda_val = value_of(lambda);
 
  auto check_lambda_not_nan
      = check_cl(function, "lambda", lambda_val, "not NaN");
  auto lambda_not_nan = !isnan(lambda_val);
  auto check_theta_bounded
      = check_cl(function, "theta", theta_val, "in the interval[0, 1]");
  auto theta_bounded = 0.0 <= theta_val && theta_val <= 1.0;
 
  auto theta_bc = rowwise_broadcast(theta_val);
  auto lambda_p_log_theta_expr = lambda_val + log(theta_bc);
  matrix_cl<double> lambda_p_log_theta;
  matrix_cl<double> lambda_p_log_theta_colwise_max;
  if (theta.cols() == lambda.cols()) {
    results(check_lambda_not_nan, check_theta_bounded, lambda_p_log_theta,
            lambda_p_log_theta_colwise_max)
        = expressions(lambda_not_nan, theta_bounded, lambda_p_log_theta_expr,
                      colwise_max(lambda_p_log_theta_expr));
  } else {
    results(check_lambda_not_nan, lambda_p_log_theta,
            lambda_p_log_theta_colwise_max)
        = expressions(lambda_not_nan, lambda_p_log_theta_expr,
                      colwise_max(lambda_p_log_theta_expr));
    check_theta_bounded = theta_bounded;
  }
  while (lambda_p_log_theta_colwise_max.rows() > 1) {
    lambda_p_log_theta_colwise_max
        = colwise_max(lambda_p_log_theta_colwise_max).eval();
  }
  matrix_cl<double> sum_exp = colwise_sum(exp(
      lambda_p_log_theta - colwise_broadcast(lambda_p_log_theta_colwise_max)));
  while (sum_exp.rows() > 1) {
    sum_exp = colwise_sum(sum_exp).eval();
  }
 
  auto logp_vec_expr = transpose(lambda_p_log_theta_colwise_max + log(sum_exp));
  matrix_cl<double> logp_vec;
  matrix_cl<double> logp_sum;
  results(logp_vec, logp_sum) = expressions(
      calc_if<!is_constant_all<T_theta_cl, T_lambda_cl>::value>(logp_vec_expr),
      colwise_sum(logp_vec_expr));
 
  auto ops_partials = make_partials_propagator(theta_col, lambda);
  if (!is_constant_all<T_theta_cl, T_lambda_cl>::value) {
    auto derivs_expr = exp(lambda_val - colwise_broadcast(transpose(logp_vec)));
    if (!is_constant<T_lambda_cl>::value) {
      auto lambda_deriv_expr = elt_multiply(derivs_expr, theta_bc);
      matrix_cl<double> derivs;
      matrix_cl<double> lambda_deriv;
      results(derivs, lambda_deriv)
          = expressions(calc_if<!is_constant<T_theta_cl>::value>(derivs_expr),
                        lambda_deriv_expr);
 
      partials<1>(ops_partials) = std::move(lambda_deriv);
      if (!is_constant<T_theta_cl>::value) {
        partials<0>(ops_partials) = rowwise_sum(derivs);
      }
    } else if (!is_constant<T_theta_cl>::value) {
      partials<0>(ops_partials) = rowwise_sum(derivs_expr);
    }
  }
  return ops_partials.build(sum(from_matrix_cl(logp_sum)));
}
}  // namespace math
}  // namespace stan
#endif
#endif