Skip to content

Commit 2d7fbe2

Browse files
andrjohnsandrjohns
authored
Merge pull request #2960 from stan-dev/opencl-binomial_logit_glm
Add OpenCL implementation for `binomial_logit_glm`
2 parents c90b3bf + b0b0979 commit 2d7fbe2

3 files changed

Lines changed: 360 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: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
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, typename T_x_cl,
28+
typename T_alpha_cl, typename T_beta_cl,
29+
require_all_prim_or_rev_kernel_expression_t<
30+
T_n_cl, T_N_cl, T_x_cl, T_alpha_cl, T_beta_cl>* = nullptr>
31+
return_type_t<T_x_cl, T_alpha_cl, T_beta_cl> binomial_logit_glm_lpmf(
32+
const T_n_cl& n, const T_N_cl& N, const T_x_cl& x, const T_alpha_cl& alpha,
33+
const T_beta_cl& beta) {
34+
static const char* function = "binomial_logit_glm_lpmf(OpenCL)";
35+
using T_partials_return = partials_return_t<T_x_cl, T_alpha_cl, T_beta_cl>;
36+
constexpr bool is_y_vector = !is_stan_scalar<T_n_cl>::value;
37+
constexpr bool is_alpha_vector = !is_stan_scalar<T_alpha_cl>::value;
38+
39+
const size_t N_instances
40+
= max(max_size(n, N, alpha), static_cast<size_t>(x.rows()));
41+
const size_t N_attributes = x.cols();
42+
43+
check_consistent_sizes(function, "Successes variable", n,
44+
"Population size parameter", N);
45+
check_consistent_size(function, "Successes variable", n, N_instances);
46+
check_consistent_size(function, "Population size parameter", N, N_instances);
47+
check_consistent_size(function, "Weight vector", beta, N_attributes);
48+
check_consistent_size(function, "Vector of intercepts", alpha, N_instances);
49+
50+
if (N_instances == 0 || N_attributes == 0) {
51+
return 0;
52+
}
53+
if (!include_summand<propto, T_x_cl, T_alpha_cl, T_beta_cl>::value) {
54+
return 0;
55+
}
56+
57+
auto&& x_val = value_of(x);
58+
auto&& alpha_val = value_of(alpha);
59+
auto&& beta_val = value_of(beta);
60+
61+
auto check_n_bounded
62+
= check_cl(function, "Successes variable", n, "in the interval [0, N]");
63+
auto n_bounded = 0 <= n && n <= N;
64+
auto check_N_nonnegative
65+
= check_cl(function, "Population size variable", n, "nonnegative");
66+
auto N_nonnegative = N >= 0;
67+
68+
auto theta_expr = matrix_vector_multiply(x_val, beta_val) + alpha_val;
69+
auto log_inv_logit_theta = log_inv_logit(theta_expr);
70+
auto log1m_inv_logit_theta = log1m_inv_logit(theta_expr);
71+
auto n_diff = N - n;
72+
auto logp_expr1 = elt_multiply(n, log_inv_logit_theta)
73+
+ elt_multiply(n_diff, log1m_inv_logit_theta);
74+
auto logp_expr
75+
= static_select<include_summand<propto, T_n_cl, T_N_cl>::value>(
76+
logp_expr1 + binomial_coefficient_log(N, n), logp_expr1);
77+
78+
constexpr bool need_theta_deriv
79+
= !is_constant_all<T_beta_cl, T_x_cl, T_alpha_cl>::value;
80+
auto theta_deriv_expr = n - elt_multiply(N, exp(log_inv_logit_theta));
81+
82+
constexpr bool need_theta_deriv_sum = need_theta_deriv && !is_alpha_vector;
83+
matrix_cl<double> logp_cl;
84+
matrix_cl<double> theta_deriv_cl;
85+
matrix_cl<double> theta_deriv_sum_cl;
86+
87+
results(check_n_bounded, check_N_nonnegative, logp_cl, theta_deriv_cl,
88+
theta_deriv_sum_cl)
89+
= 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") = isfinite(alpha_val);
98+
check_cl(function, "Weight vector", beta_val, "finite")
99+
= isfinite(beta_val);
100+
check_cl(function, "Matrix of independent variables", x_val, "finite")
101+
= isfinite(x_val);
102+
}
103+
104+
auto ops_partials = make_partials_propagator(x, alpha, beta);
105+
if (!is_constant_all<T_x_cl>::value) {
106+
partials<0>(ops_partials) = transpose(beta_val * transpose(theta_deriv_cl));
107+
}
108+
if (!is_constant_all<T_alpha_cl>::value) {
109+
if (is_alpha_vector) {
110+
partials<1>(ops_partials) = theta_deriv_cl;
111+
} else {
112+
forward_as<internal::broadcast_array<double>>(
113+
partials<1>(ops_partials))[0]
114+
= sum(from_matrix_cl(theta_deriv_sum_cl));
115+
}
116+
}
117+
if (!is_constant_all<T_beta_cl>::value) {
118+
// transposition of a vector can be done without copying
119+
const matrix_cl<double> theta_derivative_transpose_cl(
120+
theta_deriv_cl.buffer(), 1, theta_deriv_cl.rows());
121+
matrix_cl<double> edge3_partials_transpose_cl
122+
= theta_derivative_transpose_cl * x_val;
123+
partials<2>(ops_partials)
124+
= matrix_cl<double>(edge3_partials_transpose_cl.buffer(),
125+
edge3_partials_transpose_cl.cols(), 1);
126+
if (beta_val.rows() != 0) {
127+
edge<2>(ops_partials)
128+
.partials_.add_write_event(
129+
edge3_partials_transpose_cl.write_events().back());
130+
}
131+
}
132+
return ops_partials.build(logp);
133+
}
134+
135+
} // namespace math
136+
} // namespace stan
137+
138+
#endif
139+
#endif

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

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
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::binomial_logit_glm_lpmf;
10+
using stan::math::matrix_cl;
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, const auto& alpha,
102+
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, const auto& alpha,
107+
const auto& beta) {
108+
return stan::math::binomial_logit_glm_lpmf<true>(n, trials, x, alpha,
109+
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(binomial_logit_glm_lpmf_functor,
125+
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(binomial_logit_glm_lpmf_functor,
160+
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(binomial_logit_glm_lpmf_functor,
176+
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+
TEST(ProbDistributionsBinomialLogitGLM, opencl_matches_cpu_small_vector_alpha) {
182+
int N = 3;
183+
int M = 2;
184+
185+
std::vector<int> n{0, 1, 0};
186+
std::vector<int> trials{0, 1, 0};
187+
Eigen::MatrixXd x(N, M);
188+
x << -12, 46, -42, 24, 25, 27;
189+
Eigen::VectorXd beta(M);
190+
beta << 0.3, 2;
191+
Eigen::VectorXd alpha(N);
192+
alpha << 0.3, -0.8, 1.8;
193+
194+
stan::math::test::compare_cpu_opencl_prim_rev(binomial_logit_glm_lpmf_functor,
195+
n, trials, x, alpha, beta);
196+
stan::math::test::compare_cpu_opencl_prim_rev(
197+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
198+
}
199+
200+
TEST(ProbDistributionsBinomialLogitGLM, opencl_matches_cpu_big) {
201+
int N = 153;
202+
int M = 71;
203+
204+
std::vector<int> n(N);
205+
std::vector<int> trials(N);
206+
for (int i = 0; i < N; i++) {
207+
n[i] = Eigen::ArrayXi::Random(1, 1).abs()(0);
208+
trials[i] = n[i] + Eigen::ArrayXi::Random(1, 1).abs()(0);
209+
}
210+
Eigen::MatrixXd x = Eigen::MatrixXd::Random(N, M);
211+
Eigen::VectorXd beta = Eigen::VectorXd::Random(M);
212+
Eigen::VectorXd alpha = Eigen::VectorXd::Random(N);
213+
214+
stan::math::test::compare_cpu_opencl_prim_rev(binomial_logit_glm_lpmf_functor,
215+
n, trials, x, alpha, beta);
216+
stan::math::test::compare_cpu_opencl_prim_rev(
217+
binomial_logit_glm_lpmf_functor_propto, n, trials, x, alpha, beta);
218+
}
219+
220+
#endif

0 commit comments

Comments
 (0)