Fix discrete cdfs when one element of a vector has probability 0 or 1

stan/math/opencl/prim/bernoulli_lccdf.hpp

@@ -49,24 +49,22 @@ inline return_type_t<T_prob_cl> bernoulli_lccdf(const T_n_cl& n,
                                       theta_val, "in the interval [0, 1]");
   auto theta_bounded_expr = 0.0 <= theta_val && theta_val <= 1.0;
 
-  auto any_n_negative = colwise_max(cast<char>(n < 0));
   auto any_n_over_one = colwise_max(cast<char>(n >= 1));
-  auto P_expr = colwise_sum(log(theta_val));
-  auto deriv = elt_divide(1.0, theta_val);
+  // An element with probability 1 (n < 0) contributes nothing
+  auto cond = n < 0;
+  auto P_expr = colwise_sum(select(cond, 0.0, log(theta_val)));
+  auto deriv = select(cond, 0.0, elt_divide(1.0, theta_val));
 
-  matrix_cl<char> any_n_negative_cl;
   matrix_cl<char> any_n_over_one_cl;
   matrix_cl<double> P_cl;
   matrix_cl<double> deriv_cl;
 
-  results(check_theta_bounded, any_n_negative_cl, any_n_over_one_cl, P_cl,
-          deriv_cl)
-      = expressions(theta_bounded_expr, any_n_negative, any_n_over_one, P_expr,
+  results(check_theta_bounded, any_n_over_one_cl, P_cl, deriv_cl)
+      = expressions(theta_bounded_expr, any_n_over_one, P_expr,
                     calc_if<(is_autodiff_v<T_prob_cl>)>(deriv));
 
-  if (from_matrix_cl(any_n_negative_cl).maxCoeff()) {
-    return 0.0;
-  }
+  // An element with probability 0 (n >= 1) makes the result -inf; the
+  // gradients are treated as zero
   if (from_matrix_cl(any_n_over_one_cl).maxCoeff()) {
     return NEGATIVE_INFTY;
   }

stan/math/prim/prob/bernoulli_lccdf.hpp

@@ -9,6 +9,7 @@
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/select.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/sum.hpp>
 #include <stan/math/prim/functor/partials_propagator.hpp>
 
 namespace stan {
@@ -46,23 +47,19 @@ inline return_type_t<T_prob> bernoulli_lccdf(const T_n& n,
 
   auto ops_partials = make_partials_propagator(theta_ref);
 
-  // Explicit return for extreme values
-  // The gradients are technically ill-defined, but treated as zero
-  if (any(n_arr < 0)) {
-    return ops_partials.build(0.0);
-  } else if (any(n_arr >= 1)) {
+  // An element with probability 0 (n >= 1) makes the result -inf; the
+  // gradients are treated as zero
+  if (any(n_arr >= 1)) {
     return ops_partials.build(NEGATIVE_INFTY);
   }
 
-  size_t theta_size = math::size(theta_arr);
-  size_t n_size = math::size(n_arr);
-  double broadcast_n = theta_size == n_size ? 1 : n_size;
-
+  // Only the elements with n = 0 remain to contribute; an element with
+  // probability 1 (n < 0) contributes nothing
   if constexpr (is_autodiff_v<T_prob>) {
-    partials<0>(ops_partials) = inv(theta_arr) * broadcast_n;
+    partials<0>(ops_partials) = select(n_arr == 0, inv(theta_arr), 0.0);
   }
 
-  return ops_partials.build(sum(log(theta_arr)) * broadcast_n);
+  return ops_partials.build(sum(select(n_arr == 0, log(theta_arr), 0.0)));
 }
 
 }  // namespace math

stan/math/prim/prob/beta_binomial_lccdf.hpp

@@ -76,19 +76,18 @@ inline return_type_t<T_size1, T_size2> beta_binomial_lccdf(
   scalar_seq_view<T_beta_ref> beta_vec(beta_ref);
   size_t max_size_seq_view = max_size(n, N, alpha, beta);
 
-  // Explicit return for extreme values
-  // The gradients are technically ill-defined, but treated as zero
-  for (size_t i = 0; i < stan::math::size(n); i++) {
-    if (n_vec.val(i) < 0) {
-      return ops_partials.build(0.0);
+  // An element with probability 0 (n >= N) makes the result -inf; the
+  // gradients are treated as zero
+  for (size_t i = 0; i < max_size_seq_view; i++) {
+    if (n_vec.val(i) >= N_vec.val(i)) {
+      return ops_partials.build(negative_infinity());
     }
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit results for extreme values
-    // The gradients are technically ill-defined, but treated as neg infinity
-    if (n_vec.val(i) >= N_vec.val(i)) {
-      return ops_partials.build(negative_infinity());
+    // An element with probability 1 (n < 0) contributes nothing
+    if (n_vec.val(i) < 0) {
+      continue;
     }
 
     const T_partials_return n_dbl = n_vec.val(i);

stan/math/prim/prob/beta_neg_binomial_cdf.hpp

@@ -82,10 +82,10 @@ inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_cdf(
   T_partials_return cdf(1.0);
   auto ops_partials = make_partials_propagator(r_ref, alpha_ref, beta_ref);
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit return for extreme values
-    // The gradients are technically ill-defined, but treated as zero
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing
     if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return 1.0;
+      continue;
     }
     auto n_dbl = n_vec.val(i);
     auto r_dbl = r_vec.val(i);

stan/math/prim/prob/beta_neg_binomial_lccdf.hpp

@@ -68,22 +68,21 @@ inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lccdf(
   int size_n = stan::math::size(n);
   size_t max_size_seq_view = max_size(n, r, alpha, beta);
 
-  // Explicit return for extreme values
-  // The gradients are technically ill-defined, but treated as zero
+  // An element with probability 0 (n = INT_MAX, which stands for an infinite
+  // count) makes the result -inf; the gradients are treated as zero
   for (int i = 0; i < size_n; i++) {
-    if (n_vec.val(i) < 0) {
-      return 0.0;
+    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
+      return negative_infinity();
     }
   }
 
   using T_partials_return = partials_return_t<T_n, T_r, T_alpha, T_beta>;
   T_partials_return log_ccdf(0.0);
   auto ops_partials = make_partials_propagator(r_ref, alpha_ref, beta_ref);
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit return for extreme values
-    // The gradients are technically ill-defined, but treated as zero
-    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return ops_partials.build(negative_infinity());
+    // An element with probability 1 (n < 0) contributes nothing
+    if (n_vec.val(i) < 0) {
+      continue;
     }
     auto n_dbl = n_vec.val(i);
     auto r_dbl = r_vec.val(i);

stan/math/prim/prob/beta_neg_binomial_lcdf.hpp

@@ -80,10 +80,10 @@ inline return_type_t<T_r, T_alpha, T_beta> beta_neg_binomial_lcdf(
   T_partials_return log_cdf(0.0);
   auto ops_partials = make_partials_propagator(r_ref, alpha_ref, beta_ref);
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit return for extreme values
-    // The gradients are technically ill-defined, but treated as zero
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing
     if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return 0.0;
+      continue;
     }
     auto n_dbl = n_vec.val(i);
     auto r_dbl = r_vec.val(i);

stan/math/prim/prob/binomial_lccdf.hpp

@@ -69,25 +69,21 @@ inline return_type_t<T_prob> binomial_lccdf(const T_n& n, const T_N& N,
   scalar_seq_view<T_theta_ref> theta_vec(theta_ref);
   size_t max_size_seq_view = max_size(n, N, theta);
 
-  // Explicit return for extreme values
-  // The gradients are technically ill-defined,
-  // but treated as negative infinity
-  for (size_t i = 0; i < stan::math::size(n); i++) {
-    if (n_vec.val(i) < 0) {
-      return ops_partials.build(0.0);
+  // An element with probability 0 (n >= N) makes the result -inf; the
+  // gradients are treated as zero
+  for (size_t i = 0; i < max_size_seq_view; i++) {
+    if (n_vec.val(i) >= N_vec.val(i)) {
+      return ops_partials.build(NEGATIVE_INFTY);
     }
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
+    // An element with probability 1 (n < 0) contributes nothing
+    if (n_vec.val(i) < 0) {
+      continue;
+    }
     const T_partials_return n_dbl = n_vec.val(i);
     const T_partials_return N_dbl = N_vec.val(i);
-
-    // Explicit results for extreme values
-    // The gradients are technically ill-defined, but treated as zero
-    if (n_dbl >= N_dbl) {
-      return ops_partials.build(NEGATIVE_INFTY);
-    }
-
     const T_partials_return theta_dbl = theta_vec.val(i);
     const T_partials_return Pi
         = 1.0 - inc_beta(N_dbl - n_dbl, n_dbl + 1, 1 - theta_dbl);

stan/math/prim/prob/neg_binomial_2_cdf.hpp

@@ -78,10 +78,10 @@ inline return_type_t<T_location, T_precision> neg_binomial_2_cdf(
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit results for extreme values
-    // The gradients are technically ill-defined, but treated as zero
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing
     if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return ops_partials.build(1.0);
+      continue;
     }
 
     const T_partials_return n_dbl_p1 = n_vec.val(i) + 1;

stan/math/prim/prob/neg_binomial_2_lcdf.hpp

@@ -10,6 +10,8 @@
 #include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/prob/beta_cdf_log.hpp>
 #include <cmath>
+#include <limits>
+#include <vector>
 
 namespace stan {
 namespace math {
@@ -43,10 +45,32 @@ inline return_type_t<T_location, T_precision> neg_binomial_2_lcdf(
   size_t size_n = stan::math::size(n);
   size_t size_phi_mu = max_size(mu, phi);
 
+  bool any_n_max = false;
   for (size_t i = 0; i < size_n; i++) {
     if (n_vec[i] < 0) {
       return LOG_ZERO;
     }
+    any_n_max = any_n_max || n_vec[i] == std::numeric_limits<int>::max();
+  }
+
+  if (any_n_max) {
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing. Remove such elements before the
+    // beta lcdf call: the beta lcdf gradient is not accurate when the
+    // second shape is 2^31.
+    const size_t max_size_seq_view = max_size(n, mu, phi);
+    std::vector<return_type_t<T_location, T_precision>> phi_mu_k;
+    std::vector<scalar_type_t<T_precision>> phi_k;
+    std::vector<double> np1_k;
+    for (size_t i = 0; i < max_size_seq_view; i++) {
+      if (n_vec[i] == std::numeric_limits<int>::max()) {
+        continue;
+      }
+      phi_mu_k.push_back(phi_vec[i] / (phi_vec[i] + mu_vec[i]));
+      phi_k.push_back(phi_vec[i]);
+      np1_k.push_back(n_vec[i] + 1.0);
+    }
+    return beta_cdf_log(phi_mu_k, phi_k, np1_k);
   }
 
   VectorBuilder<true, return_type_t<T_location, T_precision>, T_location,

stan/math/prim/prob/neg_binomial_cdf.hpp

@@ -76,10 +76,10 @@ inline return_type_t<T_shape, T_inv_scale> neg_binomial_cdf(
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit results for extreme values
-    // The gradients are technically ill-defined, but treated as zero
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing
     if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return ops_partials.build(1.0);
+      continue;
     }
 
     const T_partials_return n_dbl = n_vec.val(i);

stan/math/prim/prob/neg_binomial_lccdf.hpp

@@ -57,11 +57,11 @@ inline return_type_t<T_shape, T_inv_scale> neg_binomial_lccdf(
   size_t size_n_alpha = max_size(n, alpha);
   size_t max_size_seq_view = max_size(n, alpha, beta_param);
 
-  // Explicit return for extreme values
-  // The gradients are technically ill-defined, but treated as zero
+  // An element with probability 0 (n = INT_MAX, which stands for an infinite
+  // count) makes the result -inf; the gradients are treated as zero
   for (size_t i = 0; i < size_n; i++) {
-    if (n_vec.val(i) < 0) {
-      return ops_partials.build(0.0);
+    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
+      return ops_partials.build(negative_infinity());
     }
   }
 
@@ -87,10 +87,9 @@ inline return_type_t<T_shape, T_inv_scale> neg_binomial_lccdf(
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit results for extreme values
-    // The gradients are technically ill-defined, but treated as zero
-    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return ops_partials.build(negative_infinity());
+    // An element with probability 1 (n < 0) contributes nothing
+    if (n_vec.val(i) < 0) {
+      continue;
     }
 
     const T_partials_return n_dbl = n_vec.val(i);

stan/math/prim/prob/neg_binomial_lcdf.hpp

@@ -74,7 +74,7 @@ inline return_type_t<T_shape, T_inv_scale> neg_binomial_lcdf(
 
   if constexpr (is_autodiff_v<T_shape>) {
     for (size_t i = 0; i < size_n; i++) {
-      digammaN_vec[i] = digamma(n_vec.val(i) + 1);
+      digammaN_vec[i] = digamma(n_vec.val(i) + 1.0);
     }
     for (size_t i = 0; i < size_alpha; i++) {
       digammaAlpha_vec[i] = digamma(alpha_vec.val(i));
@@ -87,10 +87,10 @@ inline return_type_t<T_shape, T_inv_scale> neg_binomial_lcdf(
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    // Explicit results for extreme values
-    // The gradients are technically ill-defined, but treated as zero
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing
     if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return ops_partials.build(0.0);
+      continue;
     }
 
     const T_partials_return n_dbl = n_vec.val(i);

stan/math/prim/prob/poisson_binomial_cdf.hpp

@@ -49,6 +49,11 @@ inline return_type_t<T_theta> poisson_binomial_cdf(const T_y& y,
 
   return_type_t<T_theta> lcdf = 0.0;
   for (size_t i = 0; i < max_sz; ++i) {
+    // An element with probability 1 (y equal to the number of trials)
+    // contributes nothing
+    if (y_vec[i] == static_cast<int>(stan::math::size(theta_vec[i]))) {
+      continue;
+    }
     auto x = log_sum_exp(poisson_binomial_log_probs(y_vec[i], theta_vec[i]));
     lcdf += x;
   }

stan/math/prim/prob/poisson_binomial_lccdf.hpp

@@ -3,6 +3,7 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/log_sum_exp.hpp>
 #include <stan/math/prim/fun/log1m_exp.hpp>
@@ -50,14 +51,18 @@ inline return_type_t<T_theta> poisson_binomial_lccdf(const T_y& y,
                   1.0);
   }
 
+  // An element with probability 0 (y equal to the number of trials) makes the
+  // result -inf; the gradients are treated as zero
+  for (size_t i = 0; i < max_sz; ++i) {
+    if (y_vec[i] == static_cast<int>(stan::math::size(theta_vec[i]))) {
+      return negative_infinity();
+    }
+  }
+
   return_type_t<T_theta> lccdf = 0.0;
   for (size_t i = 0; i < max_sz; ++i) {
     if (stan::math::size(theta_vec[i]) == 1) {
-      if (y_vec[i] == 0) {
-        lccdf += log(theta_vec[i][0]);
-      } else {
-        lccdf -= stan::math::INFTY;
-      }
+      lccdf += log(theta_vec[i][0]);
     } else {
       auto x = log1m_exp(
           log_sum_exp(poisson_binomial_log_probs(y_vec[i], theta_vec[i])));

stan/math/prim/prob/poisson_binomial_lcdf.hpp

@@ -50,6 +50,11 @@ inline return_type_t<T_theta> poisson_binomial_lcdf(const T_y& y,
 
   return_type_t<T_theta> lcdf = 0.0;
   for (size_t i = 0; i < max_sz; ++i) {
+    // An element with probability 1 (y equal to the number of trials)
+    // contributes nothing
+    if (y_vec[i] == static_cast<int>(stan::math::size(theta_vec[i]))) {
+      continue;
+    }
     auto x = log_sum_exp(poisson_binomial_log_probs(y_vec[i], theta_vec[i]));
     lcdf += x;
   }

stan/math/prim/prob/yule_simon_cdf.hpp

@@ -52,20 +52,22 @@ inline return_type_t<T_alpha> yule_simon_cdf(const T_n& n,
   scalar_seq_view<T_alpha_ref> alpha_vec(alpha_ref);
   const size_t max_size_seq_view = max_size(n_ref, alpha_ref);
 
-  // Explicit return for invalid or extreme values
-  // The gradients are technically ill-defined, but treated as zero
+  // An element with probability 0 (n < 1) makes the result 0; the gradients
+  // are treated as zero
   for (int i = 0; i < stan::math::size(n); i++) {
     if (n_vec.val(i) < 1.0) {
       return 0.0;
     }
-    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
-      return 1.0;
-    }
   }
 
   T_partials_return cdf(1.0);
   auto ops_partials = make_partials_propagator(alpha_ref);
   for (size_t i = 0; i < max_size_seq_view; i++) {
+    // An element with probability 1 (n = INT_MAX, which stands for an
+    // infinite count) contributes nothing
+    if (n_vec.val(i) == std::numeric_limits<int>::max()) {
+      continue;
+    }
     auto np1 = n_vec.val(i) + 1.0;
     auto ap1 = alpha_vec.val(i) + 1.0;
     auto nap1 = n_vec.val(i) + ap1;