Skip to content

Commit a9cc1bc

Browse files
andrjohnsandrjohns
committed
Add opencl implementation for binomial_logit_glm
1 parent c90b3bf commit a9cc1bc

3 files changed

Lines changed: 364 additions & 0 deletions

File tree

‎stan/math/opencl/prim.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@
112112
#include <stan/math/opencl/prim/beta_proportion_lpdf.hpp>
113113
#include <stan/math/opencl/prim/binomial_logit_lpmf.hpp>
114114
#include <stan/math/opencl/prim/binomial_lpmf.hpp>
115+
#include <stan/math/opencl/prim/binomial_logit_glm_lpmf.hpp>
115116
#include <stan/math/opencl/prim/block.hpp>
116117
#include <stan/math/opencl/prim/categorical_logit_glm_lpmf.hpp>
117118
#include <stan/math/opencl/prim/cauchy_cdf.hpp>

‎stan/math/opencl/prim/binomial_logit_glm_lpmf.hpp‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
#ifndef STAN_MATH_OPENCL_PRIM_BINOMIAL_LOGIT_GLM_LPMF_HPP
2+
#define STAN_MATH_OPENCL_PRIM_BINOMIAL_LOGIT_GLM_LPMF_HPP
3+
#ifdef STAN_OPENCL
4+
5+
#include <stan/math/opencl/prim/size.hpp>
6+
#include <stan/math/opencl/rev/operands_and_partials.hpp>
7+
#include <stan/math/opencl/copy.hpp>
8+
#include <stan/math/opencl/prim/multiply.hpp>
9+
#include <stan/math/opencl/kernel_generator.hpp>
10+
#include <stan/math/prim/meta.hpp>
11+
#include <stan/math/prim/err.hpp>
12+
#include <stan/math/prim/fun/Eigen.hpp>
13+
#include <stan/math/prim/fun/constants.hpp>
14+
#include <stan/math/prim/fun/max.hpp>
15+
#include <stan/math/prim/fun/size.hpp>
16+
#include <stan/math/prim/fun/size_zero.hpp>
17+
#include <stan/math/prim/fun/sum.hpp>
18+
#include <stan/math/prim/fun/to_ref.hpp>
19+
#include <stan/math/prim/fun/value_of.hpp>
20+
#include <stan/math/prim/fun/value_of_rec.hpp>
21+
22+
#include <cmath>
23+
24+
namespace stan {
25+
namespace math {
26+
27+
template <bool propto, typename T_n_cl, typename T_N_cl,
28+
typename T_x_cl, typename T_alpha_cl,
29+
typename T_beta_cl,
30+
require_all_prim_or_rev_kernel_expression_t<
31+
T_n_cl, T_N_cl, T_x_cl, T_alpha_cl, T_beta_cl>* = nullptr>
32+
return_type_t<T_x_cl, T_alpha_cl, T_beta_cl> binomial_logit_glm_lpmf(
33+
const T_n_cl& n, const T_N_cl& N, const T_x_cl& x, const T_alpha_cl& alpha,
34+
const T_beta_cl& beta) {
35+
static const char* function = "binomial_logit_glm_lpmf(OpenCL)";
36+
using T_partials_return = partials_return_t<T_x_cl, T_alpha_cl, T_beta_cl>;
37+
constexpr bool is_y_vector = !is_stan_scalar<T_n_cl>::value;
38+
constexpr bool is_alpha_vector = !is_stan_scalar<T_alpha_cl>::value;
39+
40+
const size_t N_instances = max(max_size(n, N, alpha),
41+
static_cast<size_t>(x.rows()));
42+
const size_t N_attributes = x.cols();
43+
44+
check_consistent_sizes(function, "Successes variable", n,
45+
"Population size parameter", N);
46+
check_consistent_size(function, "Successes variable", n, N_instances);
47+
check_consistent_size(function, "Population size parameter", N, N_instances);
48+
check_consistent_size(function, "Weight vector", beta, N_attributes);
49+
check_consistent_size(function, "Vector of intercepts", alpha, N_instances);
50+
51+
if (N_instances == 0 || N_attributes == 0) {
52+
return 0;
53+
}
54+
if (!include_summand<propto, T_x_cl, T_alpha_cl, T_beta_cl>::value) {
55+
return 0;
56+
}
57+
58+
auto&& x_val = value_of(x);
59+
auto&& alpha_val = value_of(alpha);
60+
auto&& beta_val = value_of(beta);
61+
62+
auto check_n_bounded
63+
= check_cl(function, "Successes variable", n, "in the interval [0, N]");
64+
auto n_bounded = 0 <= n && n <= N;
65+
auto check_N_nonnegative
66+
= check_cl(function, "Population size variable", n, "nonnegative");
67+
auto N_nonnegative = N >= 0;
68+
69+
auto theta_expr = matrix_vector_multiply(x_val, beta_val) + alpha_val;
70+
auto log_inv_logit_theta = log_inv_logit(theta_expr);
71+
auto log1m_inv_logit_theta = log1m_inv_logit(theta_expr);
72+
auto n_diff = N - n;
73+
auto logp_expr1 = elt_multiply(n, log_inv_logit_theta)
74+
+ elt_multiply(n_diff, log1m_inv_logit_theta);
75+
auto logp_expr
76+
= static_select<include_summand<propto, T_n_cl, T_N_cl>::value>(
77+
logp_expr1 + binomial_coefficient_log(N, n), logp_expr1);
78+
79+
constexpr bool need_theta_deriv =
80+
!is_constant_all<T_beta_cl, T_x_cl, T_alpha_cl>::value;
81+
auto theta_deriv_expr = n - elt_multiply(N, exp(log_inv_logit_theta));
82+
83+
constexpr bool need_theta_deriv_sum = need_theta_deriv && !is_alpha_vector;
84+
matrix_cl<double> logp_cl;
85+
matrix_cl<double> theta_deriv_cl;
86+
matrix_cl<double> theta_deriv_sum_cl;
87+
88+
results(check_n_bounded, check_N_nonnegative, logp_cl,
89+
theta_deriv_cl, theta_deriv_sum_cl) = expressions(
90+
n_bounded, N_nonnegative, logp_expr,
91+
calc_if<need_theta_deriv>(theta_deriv_expr),
92+
calc_if<need_theta_deriv_sum>(colwise_sum(theta_deriv_expr)));
93+
94+
T_partials_return logp = sum(from_matrix_cl(logp_cl));
95+
using std::isfinite;
96+
if (!isfinite(logp)) {
97+
check_cl(function, "Intercept", alpha_val, "finite")
98+
= isfinite(alpha_val);
99+
check_cl(function, "Weight vector", beta_val, "finite")
100+
= isfinite(beta_val);
101+
check_cl(function, "Matrix of independent variables", x_val, "finite")
102+
= isfinite(x_val);
103+
}
104+
105+
auto ops_partials = make_partials_propagator(x, alpha, beta);
106+
if (!is_constant_all<T_x_cl>::value) {
107+
partials<0>(ops_partials)
108+
= transpose(beta_val * transpose(theta_deriv_cl));
109+
}
110+
if (!is_constant_all<T_alpha_cl>::value) {
111+
if (is_alpha_vector) {
112+
partials<1>(ops_partials) = theta_deriv_cl;
113+
} else {
114+
forward_as<internal::broadcast_array<double>>(
115+
partials<1>(ops_partials))[0]
116+
= sum(from_matrix_cl(theta_deriv_sum_cl));
117+
}
118+
}
119+
if (!is_constant_all<T_beta_cl>::value) {
120+
// transposition of a vector can be done without copying
121+
const matrix_cl<double> theta_derivative_transpose_cl(
122+
theta_deriv_cl.buffer(), 1, theta_deriv_cl.rows());
123+
matrix_cl<double> edge3_partials_transpose_cl
124+
= theta_derivative_transpose_cl * x_val;
125+
partials<2>(ops_partials)
126+
= matrix_cl<double>(edge3_partials_transpose_cl.buffer(),
127+
edge3_partials_transpose_cl.cols(), 1);
128+
if (beta_val.rows() != 0) {
129+
edge<2>(ops_partials)
130+
.partials_.add_write_event(
131+
edge3_partials_transpose_cl.write_events().back());
132+
}
133+
}
134+
return ops_partials.build(logp);
135+
}
136+
137+
} // namespace math
138+
} // namespace stan
139+
140+
#endif
141+
#endif

‎test/unit/math/opencl/rev/binomial_logit_glm_lpmf_test.cpp‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
#ifdef STAN_OPENCL
2+
#include <stan/math.hpp>
3+
#include <stan/math/opencl/rev.hpp>
4+
#include <gtest/gtest.h>
5+
#include <test/unit/math/opencl/util.hpp>
6+
#include <vector>
7+
8+
TEST(ProbDistributionsBinomialLogitGLM, error_checking) {
9+
using stan::math::matrix_cl;
10+
using stan::math::binomial_logit_glm_lpmf;
11+
12+
int N = 3;
13+
int M = 2;
14+
15+
std::vector<int> n{1, 0, 1};
16+
std::vector<int> n_size{1, 0, 1, 0};
17+
std::vector<int> n_value{0, 1, -23};
18+
19+
std::vector<int> trials{10, 5, 2};
20+
std::vector<int> trials_size{1, 0, 1, 0};
21+
std::vector<int> trials_value{5, 1, -1};
22+
23+
Eigen::MatrixXd x(N, M);
24+
x << -12, 46, -42, 24, 25, 27;
25+
Eigen::MatrixXd x_size1(N - 1, M);
26+
x_size1 << -12, 46, -42, 24;
27+
Eigen::MatrixXd x_size2(N, M - 1);
28+
x_size2 << -12, 46, -42;
29+
Eigen::MatrixXd x_value(N, M);
30+
x_value << -12, 46, -42, 24, 25, -INFINITY;
31+
Eigen::VectorXd beta(M);
32+
beta << 0.3, 2;
33+
Eigen::VectorXd beta_size(M + 1);
34+
beta_size << 0.3, 2, 0.4;
35+
Eigen::VectorXd beta_value(M);
36+
beta_value << 0.3, INFINITY;
37+
Eigen::VectorXd alpha(N);
38+
alpha << 0.3, -0.8, 1.8;
39+
Eigen::VectorXd alpha_size(N - 1);
40+
alpha_size << 0.3, -0.8;
41+
Eigen::VectorXd alpha_value(N);
42+
alpha_value << 0.3, -0.8, NAN;
43+
44+
matrix_cl<double> x_cl(x);
45+
matrix_cl<double> x_size1_cl(x_size1);
46+
matrix_cl<double> x_size2_cl(x_size2);
47+
matrix_cl<double> x_value_cl(x_value);
48+
matrix_cl<int> n_cl(n);
49+
matrix_cl<int> n_size_cl(n_size);
50+
matrix_cl<int> n_value_cl(n_value);
51+
matrix_cl<int> trials_cl(trials);
52+
matrix_cl<int> trials_size_cl(trials_size);
53+
matrix_cl<int> trials_value_cl(trials_value);
54+
matrix_cl<double> beta_cl(beta);
55+
matrix_cl<double> beta_size_cl(beta_size);
56+
matrix_cl<double> beta_value_cl(beta_value);
57+
matrix_cl<double> alpha_cl(alpha);
58+
matrix_cl<double> alpha_size_cl(alpha_size);
59+
matrix_cl<double> alpha_value_cl(alpha_value);
60+
61+
EXPECT_NO_THROW(
62+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_cl, alpha_cl, beta_cl));
63+
64+
EXPECT_THROW(
65+
binomial_logit_glm_lpmf(n_size_cl, trials_cl, x_cl, alpha_cl, beta_cl),
66+
std::invalid_argument);
67+
EXPECT_THROW(
68+
binomial_logit_glm_lpmf(n_cl, trials_size_cl, x_cl, alpha_cl, beta_cl),
69+
std::invalid_argument);
70+
EXPECT_THROW(
71+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_size1_cl, alpha_cl, beta_cl),
72+
std::invalid_argument);
73+
EXPECT_THROW(
74+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_size2_cl, alpha_cl, beta_cl),
75+
std::invalid_argument);
76+
EXPECT_THROW(
77+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_cl, alpha_size_cl, beta_cl),
78+
std::invalid_argument);
79+
EXPECT_THROW(
80+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_cl, alpha_cl, beta_size_cl),
81+
std::invalid_argument);
82+
83+
EXPECT_THROW(
84+
binomial_logit_glm_lpmf(n_value_cl, trials_cl, x_cl, alpha_cl, beta_cl),
85+
std::domain_error);
86+
EXPECT_THROW(
87+
binomial_logit_glm_lpmf(n_cl, trials_value_cl, x_cl, alpha_cl, beta_cl),
88+
std::domain_error);
89+
EXPECT_THROW(
90+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_value_cl, alpha_cl, beta_cl),
91+
std::domain_error);
92+
EXPECT_THROW(
93+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_cl, alpha_value_cl, beta_cl),
94+
std::domain_error);
95+
EXPECT_THROW(
96+
binomial_logit_glm_lpmf(n_cl, trials_cl, x_cl, alpha_cl, beta_value_cl),
97+
std::domain_error);
98+
}
99+
100+
auto binomial_logit_glm_lpmf_functor
101+
= [](const auto& n, const auto& trials, const auto& x,
102+
const auto& alpha, const auto& beta) {
103+
return stan::math::binomial_logit_glm_lpmf(n, trials, x, alpha, beta);
104+
};
105+
auto binomial_logit_glm_lpmf_functor_propto
106+
= [](const auto& n, const auto& trials, const auto& x,
107+
const auto& alpha, const auto& beta) {
108+
return stan::math::binomial_logit_glm_lpmf<true>(n, trials, x,
109+
alpha, beta);
110+
};
111+
112+
TEST(ProbDistributionsBinomialLogitGLM, opencl_matches_cpu_small_simple) {
113+
int N = 3;
114+
int M = 2;
115+
116+
std::vector<int> n{0, 1, 0};
117+
std::vector<int> trials{10, 4, 15};
118+
Eigen::MatrixXd x(N, M);
119+
x << -12, 46, -42, 24, 25, 27;
120+
Eigen::VectorXd beta(M);
121+
beta << 0.3, 2;
122+
double alpha = 0.3;
123+
124+
stan::math::test::compare_cpu_opencl_prim_rev(
125+
binomial_logit_glm_lpmf_functor, n, trials, x, alpha, beta);
126+
stan::math::test::compare_cpu_opencl_prim_rev(
127+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
128+
}
129+
130+
TEST(ProbDistributionsBinomialLogitGLM, opencl_broadcast_n) {
131+
int N = 3;
132+
int M = 2;
133+
134+
int n_scal = 1;
135+
std::vector<int> trials{10, 4, 15};
136+
Eigen::MatrixXd x(N, M);
137+
x << -12, 46, -42, 24, 25, 27;
138+
Eigen::VectorXd beta(M);
139+
beta << 0.3, 2;
140+
double alpha = 0.3;
141+
142+
stan::math::test::test_opencl_broadcasting_prim_rev<0>(
143+
binomial_logit_glm_lpmf_functor, n_scal, trials, x, alpha, beta);
144+
stan::math::test::test_opencl_broadcasting_prim_rev<0>(
145+
binomial_logit_glm_lpmf_functor_propto, n_scal, trials, x, alpha, beta);
146+
}
147+
148+
TEST(ProbDistributionsBinomialLogitGLM, opencl_matches_cpu_zero_instances) {
149+
int N = 0;
150+
int M = 2;
151+
152+
std::vector<int> n{};
153+
std::vector<int> trials{};
154+
Eigen::MatrixXd x(N, M);
155+
Eigen::VectorXd beta(M);
156+
beta << 0.3, 2;
157+
double alpha = 0.3;
158+
159+
stan::math::test::compare_cpu_opencl_prim_rev(
160+
binomial_logit_glm_lpmf_functor, n, trials, x, alpha, beta);
161+
stan::math::test::compare_cpu_opencl_prim_rev(
162+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
163+
}
164+
165+
TEST(ProbDistributionsBinomialLogitGLM, opencl_matches_cpu_zero_attributes) {
166+
int N = 3;
167+
int M = 0;
168+
169+
std::vector<int> n{0, 1, 0};
170+
std::vector<int> trials{10, 5, 4};
171+
Eigen::MatrixXd x(N, M);
172+
Eigen::VectorXd beta(M);
173+
double alpha = 0.3;
174+
175+
stan::math::test::compare_cpu_opencl_prim_rev(
176+
binomial_logit_glm_lpmf_functor, n, trials, x, alpha, beta);
177+
stan::math::test::compare_cpu_opencl_prim_rev(
178+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
179+
}
180+
181+
182+
TEST(ProbDistributionsBinomialLogitGLM,
183+
opencl_matches_cpu_small_vector_alpha) {
184+
int N = 3;
185+
int M = 2;
186+
187+
std::vector<int> n{0, 1, 0};
188+
std::vector<int> trials{0, 1, 0};
189+
Eigen::MatrixXd x(N, M);
190+
x << -12, 46, -42, 24, 25, 27;
191+
Eigen::VectorXd beta(M);
192+
beta << 0.3, 2;
193+
Eigen::VectorXd alpha(N);
194+
alpha << 0.3, -0.8, 1.8;
195+
196+
stan::math::test::compare_cpu_opencl_prim_rev(
197+
binomial_logit_glm_lpmf_functor, n, trials, x, alpha, beta);
198+
stan::math::test::compare_cpu_opencl_prim_rev(
199+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
200+
}
201+
202+
TEST(ProbDistributionsBinomialLogitGLM, opencl_matches_cpu_big) {
203+
int N = 153;
204+
int M = 71;
205+
206+
std::vector<int> n(N);
207+
std::vector<int> trials(N);
208+
for (int i = 0; i < N; i++) {
209+
n[i] = Eigen::ArrayXi::Random(1, 1).abs()(0);
210+
trials[i] = n[i] + Eigen::ArrayXi::Random(1, 1).abs()(0);
211+
}
212+
Eigen::MatrixXd x = Eigen::MatrixXd::Random(N, M);
213+
Eigen::VectorXd beta = Eigen::VectorXd::Random(M);
214+
Eigen::VectorXd alpha = Eigen::VectorXd::Random(N);
215+
216+
stan::math::test::compare_cpu_opencl_prim_rev(
217+
binomial_logit_glm_lpmf_functor, n, trials, x, alpha, beta);
218+
stan::math::test::compare_cpu_opencl_prim_rev(
219+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
220+
}
221+
222+
#endif

0 commit comments

Comments
 (0)