Compute GLM design-matrix products with multiply to avoid per-coefficient autodiff nodes in forward mode

stan/math/prim/prob/bernoulli_logit_glm_lpmf.hpp

@@ -4,9 +4,11 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
 #include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/isfinite.hpp>
 #include <stan/math/prim/fun/size.hpp>
@@ -106,7 +108,7 @@ inline return_type_t<T_x, T_alpha, T_beta> bernoulli_logit_glm_lpmf(
     T_ytheta_tmp ytheta_tmp = (x_val * beta_val_vec)(0, 0);
     ytheta = signs * (ytheta_tmp + as_array_or_scalar(alpha_val_vec));
   } else {
-    ytheta = (x_val * beta_val_vec).array();
+    ytheta = multiply(x_val, beta_val_vec).array();
     ytheta = signs * (ytheta + as_array_or_scalar(alpha_val_vec));
   }
 
@@ -138,8 +140,11 @@ inline return_type_t<T_x, T_alpha, T_beta> bernoulli_logit_glm_lpmf(
     if constexpr (is_autodiff_v<T_beta>) {
       if constexpr (T_x_rows == 1) {
         edge<2>(ops_partials).partials_ = theta_derivative.sum() * x_val;
+      } else if constexpr (T_x::ColsAtCompileTime == 1) {
+        partials<2>(ops_partials)[0] = dot_product(x_val, theta_derivative);
       } else {
-        partials<2>(ops_partials) = x_val.transpose() * theta_derivative;
+        partials<2>(ops_partials)
+            = multiply(x_val.transpose(), theta_derivative);
       }
     }
     if constexpr (is_autodiff_v<T_x>) {

stan/math/prim/prob/binomial_logit_glm_lpmf.hpp

@@ -3,7 +3,9 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/isfinite.hpp>
 #include <stan/math/prim/fun/size.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
@@ -102,7 +104,7 @@ inline return_type_t<T_x, T_alpha, T_beta> binomial_logit_glm_lpmf(
   if constexpr (T_x_rows == 1) {
     theta = (x_val * beta_val)(0, 0) + alpha_val;
   } else {
-    theta = (x_val * beta_val).array() + alpha_val;
+    theta = multiply(x_val, beta_val).array() + alpha_val;
   }
 
   constexpr bool gradients_calc = is_any_autodiff_v<T_beta, T_x, T_alpha>;
@@ -131,8 +133,11 @@ inline return_type_t<T_x, T_alpha, T_beta> binomial_logit_glm_lpmf(
     if constexpr (is_autodiff_v<T_beta>) {
       if constexpr (T_x_rows == 1) {
         edge<2>(ops_partials).partials_ = theta_derivative.sum() * x_val;
+      } else if constexpr (T_x::ColsAtCompileTime == 1) {
+        partials<2>(ops_partials)[0] = dot_product(x_val, theta_derivative);
       } else {
-        partials<2>(ops_partials) = x_val.transpose() * theta_derivative;
+        partials<2>(ops_partials)
+            = multiply(x_val.transpose(), theta_derivative);
       }
     }
 

stan/math/prim/prob/categorical_logit_glm_lpmf.hpp

@@ -3,6 +3,8 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
 #include <stan/math/prim/fun/exp.hpp>
@@ -91,8 +93,12 @@ inline return_type_t<T_x, T_alpha, T_beta> categorical_logit_glm_lpmf(
 
   const auto& alpha_val_vec = as_column_vector_or_scalar(alpha_val).transpose();
 
-  Array<T_partials_return, T_x_rows, Dynamic> lin
-      = (x_val * beta_val).rowwise() + alpha_val_vec;
+  Array<T_partials_return, T_x_rows, Dynamic> lin;
+  if constexpr (T_x_rows == 1 && T_beta::ColsAtCompileTime == 1) {
+    lin = dot_product(x_val, beta_val) + alpha_val_vec.array();
+  } else {
+    lin = multiply(x_val, beta_val).rowwise() + alpha_val_vec;
+  }
   Array<T_partials_return, T_x_rows, 1> lin_max
       = lin.rowwise().maxCoeff();  // This is used to prevent overflow when
                                    // calculating softmax/log_sum_exp and
@@ -133,10 +139,19 @@ inline return_type_t<T_x, T_alpha, T_beta> categorical_logit_glm_lpmf(
       for (int i = 1; i < N_instances; i++) {
         beta_y += beta_val.col(y_seq[i] - 1).array();
       }
-      edge<0>(ops_partials).partials_
-          = beta_y
-            - (exp_lin.matrix() * beta_val.transpose()).array().colwise()
-                  * inv_sum_exp_lin * N_instances;
+      if constexpr (T_beta::RowsAtCompileTime == 1) {
+        edge<0>(ops_partials).partials_
+            = beta_y
+              - dot_product(exp_lin.matrix(), beta_val) * inv_sum_exp_lin
+                    * N_instances;
+      } else {
+        edge<0>(ops_partials).partials_
+            = beta_y
+              - multiply(exp_lin.matrix(), beta_val.transpose())
+                        .array()
+                        .colwise()
+                    * inv_sum_exp_lin * N_instances;
+      }
     } else {
       Array<T_beta_partials, Dynamic, Dynamic> beta_y(N_instances,
                                                       N_attributes);
@@ -145,7 +160,7 @@ inline return_type_t<T_x, T_alpha, T_beta> categorical_logit_glm_lpmf(
       }
       edge<0>(ops_partials).partials_
           = beta_y
-            - (exp_lin.matrix() * beta_val.transpose()).array().colwise()
+            - multiply(exp_lin.matrix(), beta_val.transpose()).array().colwise()
                   * inv_sum_exp_lin;
       // TODO(Tadej) maybe we can replace previous block with the following
       // line when we have newer Eigen  partials<0>(ops_partials) = beta_val(y
@@ -169,8 +184,7 @@ inline return_type_t<T_x, T_alpha, T_beta> categorical_logit_glm_lpmf(
     }
     if constexpr (is_autodiff_v<T_beta>) {
       Matrix<T_partials_return, Dynamic, Dynamic> beta_derivative
-          = x_val.transpose().template cast<T_partials_return>()
-            * neg_softmax_lin.matrix();
+          = multiply(x_val.transpose(), neg_softmax_lin.matrix());
       if constexpr (T_x_rows == 1) {
         beta_derivative *= N_instances;
       }

stan/math/prim/prob/neg_binomial_2_log_glm_lpmf.hpp

@@ -3,10 +3,12 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/digamma.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/lgamma.hpp>
 #include <stan/math/prim/fun/log.hpp>
@@ -142,7 +144,7 @@ neg_binomial_2_log_glm_lpmf(const T_y& y, const T_x& x, const T_alpha& alpha,
     T_theta_tmp theta_tmp = (x_val * beta_val_vec)(0, 0);
     theta = theta_tmp + as_array_or_scalar(alpha_val_vec);
   } else {
-    theta = (x_val * beta_val_vec).array();
+    theta = multiply(x_val, beta_val_vec).array();
     theta += as_array_or_scalar(alpha_val_vec);
   }
   check_finite(function, "Matrix of independent variables", theta);
@@ -197,9 +199,12 @@ neg_binomial_2_log_glm_lpmf(const T_y& y, const T_x& x, const T_alpha& alpha,
       if constexpr (is_autodiff_v<T_beta>) {
         if constexpr (T_x_rows == 1) {
           edge<2>(ops_partials).partials_ = theta_derivative.sum() * x_val;
+        } else if constexpr (T_x::ColsAtCompileTime == 1) {
+          edge<2>(ops_partials).partials_[0]
+              = dot_product(x_val, theta_derivative);
         } else {
           edge<2>(ops_partials).partials_
-              = x_val.transpose() * theta_derivative;
+              = multiply(x_val.transpose(), theta_derivative);
         }
       }
       if constexpr (is_autodiff_v<T_x>) {

stan/math/prim/prob/normal_id_glm_lpdf.hpp

@@ -3,9 +3,11 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
 #include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/isfinite.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/size.hpp>
@@ -127,7 +129,7 @@ inline return_type_t<T_y, T_x, T_alpha, T_beta, T_scale> normal_id_glm_lpdf(
                 - as_array_or_scalar(alpha_val_vec))
                * inv_sigma;
   } else {
-    y_scaled = x_val * beta_val_vec;
+    y_scaled = multiply(x_val, beta_val_vec);
     y_scaled = (as_array_or_scalar(y_val_vec) - y_scaled
                 - as_array_or_scalar(alpha_val_vec))
                * inv_sigma;
@@ -156,8 +158,10 @@ inline return_type_t<T_y, T_x, T_alpha, T_beta, T_scale> normal_id_glm_lpdf(
     if constexpr (is_autodiff_v<T_beta>) {
       if constexpr (T_x_rows == 1) {
         edge<3>(ops_partials).partials_ = mu_derivative.sum() * x_val;
+      } else if constexpr (T_x::ColsAtCompileTime == 1) {
+        partials<3>(ops_partials)[0] = dot_product(mu_derivative, x_val);
       } else {
-        partials<3>(ops_partials) = mu_derivative.transpose() * x_val;
+        partials<3>(ops_partials) = multiply(mu_derivative.transpose(), x_val);
       }
     }
     if constexpr (is_autodiff_v<T_alpha>) {

stan/math/prim/prob/ordered_logistic_glm_lpmf.hpp

@@ -3,8 +3,10 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/isfinite.hpp>
 #include <stan/math/prim/fun/log1m_exp.hpp>
@@ -120,7 +122,7 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
     }
   }
 
-  T_location location = x_val * beta_val_vec;
+  T_location location = multiply(x_val, beta_val_vec);
   if (!isfinite(sum(location))) {
     check_finite(function, "Weight vector", beta);
     check_finite(function, "Matrix of independent variables", x);
@@ -189,9 +191,12 @@ inline return_type_t<T_x, T_beta, T_cuts> ordered_logistic_glm_lpmf(
           edge<1>(ops_partials).partials_
               = (location_derivative * x_val.replicate(N_instances, 1))
                     .transpose();
+        } else if constexpr (T_x::ColsAtCompileTime == 1) {
+          edge<1>(ops_partials).partials_[0]
+              = dot_product(location_derivative, x_val);
         } else {
           edge<1>(ops_partials).partials_
-              = (location_derivative * x_val).transpose();
+              = multiply(location_derivative, x_val).transpose();
         }
       }
     }

stan/math/prim/prob/poisson_log_glm_lpmf.hpp

@@ -3,9 +3,11 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/multiply.hpp>
 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
 #include <stan/math/prim/fun/constants.hpp>
+#include <stan/math/prim/fun/dot_product.hpp>
 #include <stan/math/prim/fun/exp.hpp>
 #include <stan/math/prim/fun/isfinite.hpp>
 #include <stan/math/prim/fun/lgamma.hpp>
@@ -103,7 +105,7 @@ inline return_type_t<T_x, T_alpha, T_beta> poisson_log_glm_lpmf(
     T_theta_tmp theta_tmp = (x_val * beta_val_vec).coeff(0, 0);
     theta = theta_tmp + as_array_or_scalar(alpha_val_vec);
   } else {
-    theta = x_val * beta_val_vec;
+    theta = multiply(x_val, beta_val_vec);
     theta += as_array_or_scalar(alpha_val_vec);
   }
 
@@ -129,8 +131,10 @@ inline return_type_t<T_x, T_alpha, T_beta> poisson_log_glm_lpmf(
   if constexpr (is_autodiff_v<T_beta>) {
     if constexpr (T_x_rows == 1) {
       edge<2>(ops_partials).partials_ = theta_derivative.sum() * x_val;
+    } else if constexpr (T_x::ColsAtCompileTime == 1) {
+      partials<2>(ops_partials)[0] = dot_product(x_val, theta_derivative);
     } else {
-      partials<2>(ops_partials) = x_val.transpose() * theta_derivative;
+      partials<2>(ops_partials) = multiply(x_val.transpose(), theta_derivative);
     }
   }
   if constexpr (is_autodiff_v<T_x>) {

test/unit/math/mix/prob/bernoulli_logit_glm_lpmf_test.cpp

@@ -23,3 +23,17 @@ TEST_F(AgradRev, mathMixScalFun_bernoulli_logit_glm_lpmf) {
   stan::test::expect_ad(f(y), x_rowvec, alpha, beta);
   stan::test::expect_ad(f(y), x_rowvec, alpha[0], beta);
 }
+
+// x is a column vector: a design matrix with a single attribute
+TEST_F(AgradRev, mathMixScalFun_bernoulli_logit_glm_lpmf_x_column_vector) {
+  std::vector<int> y{0, 1};
+  auto f = [y](const auto& x, const auto& alpha, const auto& beta) {
+    return stan::math::bernoulli_logit_glm_lpmf(y, x, alpha, beta);
+  };
+
+  Eigen::VectorXd x = Eigen::VectorXd::Random(2);
+  Eigen::VectorXd alpha = Eigen::VectorXd::Random(2);
+  Eigen::VectorXd beta = Eigen::VectorXd::Random(1);
+
+  stan::test::expect_ad(f, x, alpha, beta);
+}

test/unit/math/mix/prob/binomial_logit_glm_lpmf_test.cpp

@@ -32,3 +32,18 @@ TEST(mathMixScalFun, binomial_logit_glm_lpmf) {
   stan::test::expect_ad(f(n_arr, N_arr[0]), x_rowvec, alpha[0], beta);
   stan::test::expect_ad(f(n_arr, N_arr), x_rowvec, alpha[0], beta);
 }
+
+// x is a column vector: a design matrix with a single attribute
+TEST(mathMixScalFun, binomial_logit_glm_lpmf_x_column_vector) {
+  std::vector<int> n{1, 4};
+  std::vector<int> N{10, 45};
+  auto f = [n, N](const auto& x, const auto& alpha, const auto& beta) {
+    return stan::math::binomial_logit_glm_lpmf(n, N, x, alpha, beta);
+  };
+
+  Eigen::VectorXd x = Eigen::VectorXd::Random(2);
+  Eigen::VectorXd alpha = Eigen::VectorXd::Random(2);
+  Eigen::VectorXd beta = Eigen::VectorXd::Random(1);
+
+  stan::test::expect_ad(f, x, alpha, beta);
+}

test/unit/math/mix/prob/categorical_logit_glm_lpmf_test.cpp

@@ -19,3 +19,37 @@ TEST_F(AgradRev, mathMixScalFun_categorical_logit_glm_lpmf) {
   stan::test::expect_ad(f(y), x, alpha, beta);
   stan::test::expect_ad(f(y), x_rowvec, alpha, beta);
 }
+
+// several categories: beta is a matrix, so the products inside are
+// matrix-matrix products, which must also work with forward-mode scalars
+TEST_F(AgradRev, mathMixScalFun_categorical_logit_glm_lpmf_matrix_beta) {
+  auto f = [](const auto y) {
+    return [=](const auto& x, const auto& alpha, const auto& beta) {
+      return stan::math::categorical_logit_glm_lpmf(y, x, alpha, beta);
+    };
+  };
+
+  std::vector<int> y{1, 3, 2};
+  Eigen::MatrixXd x = Eigen::MatrixXd::Random(3, 2);
+  Eigen::RowVectorXd x_rowvec = x.row(0);
+  Eigen::VectorXd alpha = Eigen::VectorXd::Random(3);
+  Eigen::MatrixXd beta = Eigen::MatrixXd::Random(2, 3);
+
+  stan::test::expect_ad(f(y[0]), x_rowvec, alpha, beta);
+  stan::test::expect_ad(f(y), x, alpha, beta);
+  stan::test::expect_ad(f(y), x_rowvec, alpha, beta);
+}
+
+// x and beta are both row vectors: a single attribute
+TEST_F(AgradRev, mathMixScalFun_categorical_logit_glm_lpmf_row_vector_beta) {
+  std::vector<int> y{1, 3};
+  auto f = [y](const auto& x, const auto& alpha, const auto& beta) {
+    return stan::math::categorical_logit_glm_lpmf(y, x, alpha, beta);
+  };
+
+  Eigen::RowVectorXd x = Eigen::RowVectorXd::Random(1);
+  Eigen::VectorXd alpha = Eigen::VectorXd::Random(3);
+  Eigen::RowVectorXd beta = Eigen::RowVectorXd::Random(3);
+
+  stan::test::expect_ad(f, x, alpha, beta);
+}