diff --git a/stan/math/prim/fun/select.hpp b/stan/math/prim/fun/select.hpp index a52e7e4954b..f5ea8814d5c 100644 --- a/stan/math/prim/fun/select.hpp +++ b/stan/math/prim/fun/select.hpp @@ -4,9 +4,20 @@ #include #include #include +#include namespace stan { namespace math { +namespace internal { +template +struct is_default_ternary : std::false_type {}; + +template +struct is_default_ternary() + : std::declval())>> + : std::true_type {}; +} // namespace internal /** * If first argument is true return the second argument, @@ -14,32 +25,35 @@ namespace math { * * `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 , - require_all_stan_scalar_t* = 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>* = nullptr> +inline auto select(const bool c, T_true&& y_true, T_false&& y_false) { + if constexpr (is_container_v || is_container_v) { + check_matching_dims("select", "left hand side", y_true, "right hand side", + y_false); + } + return c ? std::forward(y_true) : std::forward(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, typename T_true_plain = promote_scalar_t>, typename T_false_plain = promote_scalar_t>, - require_all_container_t* = nullptr, + require_not_t>* = nullptr, + require_any_t, is_container>, + conjunction, + is_stan_scalar>>* = nullptr, require_all_same_t* = 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) { + check_matching_dims("select", "left hand side", y_true, "right hand side", + y_false); + } return c ? T_true_plain(std::forward(y_true)) : T_true_plain(std::forward(y_false)); } diff --git a/test/unit/math/prim/fun/select_test.cpp b/test/unit/math/prim/fun/select_test.cpp index 1086b27e9fc..13014c847ca 100644 --- a/test/unit/math/prim/fun/select_test.cpp +++ b/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; + using VecTransT = Eigen::Transpose; + + 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); +}