Allow select to return eigen expressions if both branches match

stan/math/prim/fun/select.hpp

@@ -4,42 +4,56 @@
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/functor/apply_scalar_binary.hpp>
+#include <stan/math/prim/meta/conjunction.hpp>
 
 namespace stan {
 namespace math {
+namespace internal {
+template <typename T_true, typename T_false, typename = void>
+struct is_default_ternary : std::false_type {};
+
+template <typename T_true, typename T_false>
+struct is_default_ternary<T_true, T_false,
+                          std::void_t<decltype(true ? std::declval<T_true>()
+                                                    : std::declval<T_false>())>>
+    : std::true_type {};
+}  // namespace internal
 
 /**
  * If first argument is true return the second argument,
  * else return the third argument.
  *
  * `select(c, y1, y0) = c ? y1 : y0`.
  *
- * @tparam T_true A stan `Scalar` type
- * @tparam T_false A stan `Scalar` type
+ * @tparam T_true Type of first argument, can be any type
+ * @tparam T_false Type of second argument, can be any type
+ *   convertible with the first
  * @param c Boolean condition value.
  * @param y_true Value to return if condition is true.
  * @param y_false Value to return if condition is false.
  */
 template <typename T_true, typename T_false,
-          typename ReturnT = return_type_t<T_true, T_false>,
-          require_all_stan_scalar_t<T_true, T_false>* = nullptr>
-inline ReturnT select(const bool c, const T_true y_true,
-                      const T_false y_false) {
-  return c ? ReturnT(y_true) : ReturnT(y_false);
+          require_t<internal::is_default_ternary<T_true, T_false>>* = nullptr>
+inline auto select(const bool c, T_true&& y_true, T_false&& y_false) {
+  if constexpr (is_container_v<T_true> || is_container_v<T_false>) {
+    check_matching_dims("select", "left hand side", y_true, "right hand side",
+                        y_false);
+  }
+  return c ? std::forward<T_true>(y_true) : std::forward<T_false>(y_false);
 }
 
 /**
  * If first argument is true return the second argument,
- * else return the third argument. Eigen expressions are
- * evaluated so that the return type is the same for both branches.
+ * else return the third argument. If the two branches have
+ * different types, but matching plain types (e.g., Eigen expressions),
+ * the return type is the plain type and the expression is evaluated.
  *
- * Both containers must have the same plain type. The scalar type
+ * Both branches must have the same plain type. The scalar type
  * of the return is determined by the return_type_t<> type trait.
  *
- * Overload for use with two containers.
- *
- * @tparam T_true A container of stan `Scalar` types
- * @tparam T_false A container of stan `Scalar` types
+ * @tparam T_true Type of first argument, can be any type
+ * @tparam T_false Type of second argument, can be any type
+ *   (with the same plain type as T_true)
  * @param c Boolean condition value.
  * @param y_true Value to return if condition is true.
  * @param y_false Value to return if condition is false.
@@ -49,11 +63,16 @@ template <
     typename T_return = return_type_t<T_true, T_false>,
     typename T_true_plain = promote_scalar_t<T_return, plain_type_t<T_true>>,
     typename T_false_plain = promote_scalar_t<T_return, plain_type_t<T_false>>,
-    require_all_container_t<T_true, T_false>* = nullptr,
+    require_not_t<internal::is_default_ternary<T_true, T_false>>* = nullptr,
+    require_any_t<conjunction<is_container<T_true>, is_container<T_false>>,
+                  conjunction<is_stan_scalar<T_true>,
+                              is_stan_scalar<T_false>>>* = nullptr,
     require_all_same_t<T_true_plain, T_false_plain>* = nullptr>
 inline T_true_plain select(const bool c, T_true&& y_true, T_false&& y_false) {
-  check_matching_dims("select", "left hand side", y_true, "right hand side",
-                      y_false);
+  if constexpr (is_container_v<T_true_plain>) {
+    check_matching_dims("select", "left hand side", y_true, "right hand side",
+                        y_false);
+  }
   return c ? T_true_plain(std::forward<T_true>(y_true))
            : T_true_plain(std::forward<T_false>(y_false));
 }

test/unit/math/prim/fun/select_test.cpp

@@ -202,3 +202,23 @@ TEST(MathFunctions, select_array_bool) {
   EXPECT_THROW(select(val < 0, b_array_short, b_array_short),
                std::invalid_argument);
 }
+
+TEST(MathFunction, select_same_expression) {
+  using stan::math::rep_matrix;
+  using stan::math::select;
+  using VecReplT = Eigen::Replicate<Eigen::VectorXd, -1, -1>;
+  using VecTransT = Eigen::Transpose<Eigen::VectorXd>;
+
+  Eigen::VectorXd v(3);
+  v << 1.0, 4.0, 9.0;
+  Eigen::RowVectorXd rv(3);
+  rv << 1.0, 4.0, 9.0;
+
+  Eigen::MatrixXd rtn = select(true, rep_matrix(v, 3), rep_matrix(rv, 3));
+
+  VecReplT repl_expr = rep_matrix(v, 3);
+  VecTransT trans_expr = stan::math::transpose(v);
+
+  VecReplT replicated_v = select(true, repl_expr, repl_expr);
+  VecTransT transp_v = select(false, trans_expr, trans_expr);
+}