From b9c50a2e6704a01c0aec2b0c7366aaf7c84be9cb Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Mon, 30 Mar 2020 20:45:57 +0800 Subject: [PATCH 1/4] generalise adj_jac_apply --- stan/math/prim/fun/softmax.hpp | 17 ++-- stan/math/prim/fun/value_of.hpp | 40 ++------- stan/math/rev/fun/softmax.hpp | 36 ++++---- stan/math/rev/functor/adj_jac_apply.hpp | 115 +++++++++--------------- 4 files changed, 78 insertions(+), 130 deletions(-) diff --git a/stan/math/prim/fun/softmax.hpp b/stan/math/prim/fun/softmax.hpp index 3637c0e03a9..bfc6b8658fa 100644 --- a/stan/math/prim/fun/softmax.hpp +++ b/stan/math/prim/fun/softmax.hpp @@ -41,14 +41,15 @@ namespace math { * @param[in] v Vector to transform. * @return Unit simplex result of the softmax transform of the vector. */ -template -inline Eigen::Matrix softmax( - const Eigen::Matrix& v) { - using std::exp; - check_nonzero_size("softmax", "v", v); - Eigen::Matrix theta(v.size()); - theta = (v.array() - v.maxCoeff()).exp(); - return theta.array() / theta.sum(); +template >...> +inline auto softmax(const Container& x) { + return apply_vector_unary::apply(x, [](const auto& v) { + const Eigen::Ref>& v_ref = v; + check_nonzero_size("softmax", "v", v_ref); + plain_type_t theta(v_ref.size()); + theta = (v_ref.array() - v_ref.maxCoeff()).exp(); + return (theta.array() / theta.sum()).eval(); + }); } } // namespace math diff --git a/stan/math/prim/fun/value_of.hpp b/stan/math/prim/fun/value_of.hpp index a431a2e64ac..a86d0d02915 100644 --- a/stan/math/prim/fun/value_of.hpp +++ b/stan/math/prim/fun/value_of.hpp @@ -24,7 +24,7 @@ namespace math { * @param x scalar to convert to double * @return value of scalar cast to double */ -template +template ...> inline double value_of(const T x) { return static_cast(x); } @@ -118,36 +118,9 @@ inline const std::vector& value_of(const std::vector& x) { return x; } * @param[in] M Matrix to be converted * @return Matrix of values **/ -template -inline Eigen::Matrix::type, R, C> value_of( - const Eigen::Matrix& M) { - Eigen::Matrix::type, R, C> Md(M.rows(), M.cols()); - for (int j = 0; j < M.cols(); j++) { - for (int i = 0; i < M.rows(); i++) { - Md(i, j) = value_of(M(i, j)); - } - } - return Md; -} - -/** - * Return the specified argument. - * - *

See value_of(T) for a polymorphic - * implementation using static casts. - * - *

This inline pass-through no-op should be compiled away. - * - * @tparam R number of rows in the matrix, can be Eigen::Dynamic - * @tparam C number of columns in the matrix, can be Eigen::Dynamic - * - * @param x Specified matrix. - * @return Specified matrix. - */ -template -inline const Eigen::Matrix& value_of( - const Eigen::Matrix& x) { - return x; +template ...> +inline decltype(auto) value_of(const T& M) { + return M.unaryExpr([](const auto& x){ return value_of(x); }).eval(); } /** @@ -164,9 +137,8 @@ inline const Eigen::Matrix& value_of( * @param x Specified matrix. * @return Specified matrix. */ -template -inline const Eigen::Matrix& value_of( - const Eigen::Matrix& x) { +template ...> +inline const T& value_of(const T& x) { return x; } diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp index 77e66019acb..3096d574d62 100644 --- a/stan/math/rev/fun/softmax.hpp +++ b/stan/math/rev/fun/softmax.hpp @@ -26,17 +26,16 @@ class softmax_op { * @param alpha Unconstrained input vector. * @return Softmax of the input. */ - template - Eigen::VectorXd operator()(const std::array& /* needs_adj */, - const Eigen::VectorXd& alpha) { + template + plain_type_t operator()(const std::array& /* needs_adj */, + const T& alpha) { N_ = alpha.size(); y_ = ChainableStack::instance_->memalloc_.alloc_array(N_); + Eigen::Map> y_map(y_, N_); - auto y = softmax(alpha); - for (int n = 0; n < N_; ++n) { - y_[n] = y(n); - } - return y; + y_map = softmax(alpha); + + return y_map; } /** @@ -48,12 +47,12 @@ class softmax_op { * @param adj Eigen::VectorXd of adjoints at the output of the softmax * @return Eigen::VectorXd of adjoints propagated through softmax operation */ - template - std::tuple multiply_adjoint_jacobian( + template > + std::tuple multiply_adjoint_jacobian( const std::array& /* needs_adj */, - const Eigen::VectorXd& adj) const { - vector_d adj_times_jac(N_); - Eigen::Map y(y_, N_); + const T& adj) const { + T_plain adj_times_jac(N_); + Eigen::Map y(y_, N_); adj_times_jac = -y * adj.dot(y) + y.cwiseProduct(adj); @@ -70,11 +69,14 @@ class softmax_op { * @return Softmax of the input. * @throw std::domain_error If the input vector is size 0. */ -inline Eigen::Matrix softmax( - const Eigen::Matrix& alpha) { - check_nonzero_size("softmax", "alpha", alpha); +template >>...> +inline auto softmax(const T& x) { + return apply_vector_unary::apply(x, [&](const auto& alpha) { + const Eigen::Ref>& a_ref = alpha; + check_nonzero_size("softmax", "alpha", alpha); - return adj_jac_apply(alpha); + return adj_jac_apply(a_ref); + }); } } // namespace math diff --git a/stan/math/rev/functor/adj_jac_apply.hpp b/stan/math/rev/functor/adj_jac_apply.hpp index 26409600952..8a39b769798 100644 --- a/stan/math/rev/functor/adj_jac_apply.hpp +++ b/stan/math/rev/functor/adj_jac_apply.hpp @@ -56,13 +56,10 @@ inline void build_y_adj(vari** y_vi, const std::array& M, * @param[in] M shape of y_adj * @param[out] y_adj reference to Eigen::Matrix where adjoints are to be stored */ -template +template ...> inline void build_y_adj(vari** y_vi, const std::array& M, - Eigen::Matrix& y_adj) { - y_adj.resize(M[0], M[1]); - for (int m = 0; m < y_adj.size(); ++m) { - y_adj(m) = y_vi[m]->adj_; - } + EigT& y_adj) { + y_adj = Eigen::Map(y_vi, M[0], M[1]).adj(); } /** @@ -70,15 +67,15 @@ inline void build_y_adj(vari** y_vi, const std::array& M, * definition of dimensionality is deferred to specializations. By * default don't have a value (fail to compile) */ -template +template struct compute_dims {}; /** * Compute the dimensionality of the given template argument. Double * types have dimensionality zero. */ -template <> -struct compute_dims { +template +struct compute_dims> { static constexpr size_t value = 0; }; @@ -87,7 +84,7 @@ struct compute_dims { * std::vector has dimension 1 */ template -struct compute_dims> { +struct compute_dims> { static constexpr size_t value = 1; }; @@ -99,8 +96,8 @@ struct compute_dims> { * @tparam R number of rows, can be Eigen::Dynamic * @tparam C number of columns, can be Eigen::Dynamic */ -template -struct compute_dims> { +template +struct compute_dims> { static constexpr size_t value = 2; }; } // namespace internal @@ -121,7 +118,7 @@ template struct adj_jac_vari : public vari { std::array is_var_; using FReturnType - = std::result_of_t; + = std::result_of_t()))...)>; F f_; std::array offsets_; @@ -153,8 +150,8 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template - inline size_t count_memory(size_t count, const Eigen::Matrix& x, + template ...> + inline size_t count_memory(size_t count, const T& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -162,8 +159,9 @@ struct adj_jac_vari : public vari { return count_memory(count, args...); } - template - inline size_t count_memory(size_t count, const Eigen::Matrix& x, + template ...> + inline size_t count_memory(size_t count, const T& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -179,16 +177,9 @@ struct adj_jac_vari : public vari { return count_memory(count, args...); } - template - inline size_t count_memory(size_t count, const std::vector& x, - const Pargs&... args) { - static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; - offsets_[t] = count; - return count_memory(count, args...); - } - - template - inline size_t count_memory(size_t count, const std::vector& x, + template ...> + inline size_t count_memory(size_t count, const T& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -203,21 +194,14 @@ struct adj_jac_vari : public vari { return count_memory(count, args...); } - template - inline size_t count_memory(size_t count, const double& x, + template ...> + inline size_t count_memory(size_t count, const T& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; return count_memory(count, args...); } - template - inline size_t count_memory(size_t count, const int& x, const Pargs&... args) { - static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; - offsets_[t] = count; - return count_memory(count, args...); - } - size_t count_memory(size_t count) const { return count; } /** @@ -238,19 +222,16 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template - inline void prepare_x_vis(const Eigen::Matrix& x, - const Pargs&... args) { + template ...> + inline void prepare_x_vis(const T& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; - for (int i = 0; i < x.size(); ++i) { - x_vis_[offsets_[t] + i] = x(i).vi_; - } + Eigen::Map(&x_vis_[offsets_[t]], x.rows(), x.cols()) + = x.matrix().vi(); prepare_x_vis(args...); } - template - inline void prepare_x_vis(const Eigen::Matrix& x, - const Pargs&... args) { + template ...> + inline void prepare_x_vis(const T& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -263,17 +244,13 @@ struct adj_jac_vari : public vari { prepare_x_vis(args...); } - template - inline void prepare_x_vis(const std::vector& x, + template ...> + inline void prepare_x_vis(const T& x, const Pargs&... args) { prepare_x_vis(args...); } - template - inline void prepare_x_vis(const std::vector& x, const Pargs&... args) { - prepare_x_vis(args...); - } - template inline void prepare_x_vis(const var& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; @@ -281,13 +258,8 @@ struct adj_jac_vari : public vari { prepare_x_vis(args...); } - template - inline void prepare_x_vis(const double& x, const Pargs&... args) { - prepare_x_vis(args...); - } - - template - inline void prepare_x_vis(const int& x, const Pargs&... args) { + template ...> + inline void prepare_x_vis(const T& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -346,19 +318,20 @@ struct adj_jac_vari : public vari { * @param val_y output of F::operator() * @return Eigen::Matrix of vars */ - template - inline Eigen::Matrix build_return_varis_and_vars( - const Eigen::Matrix& val_y) { + template , + require_t>...> + inline auto build_return_varis_and_vars(const T& val_y) { + constexpr int R = T_plain::RowsAtCompileTime; + constexpr int C = T_plain::ColsAtCompileTime; M_[0] = val_y.rows(); M_[1] = val_y.cols(); Eigen::Matrix var_y(M_[0], M_[1]); y_vi_ = ChainableStack::instance_->memalloc_.alloc_array(var_y.size()); - for (int m = 0; m < var_y.size(); ++m) { - y_vi_[m] = new vari(val_y(m), false); - var_y(m) = y_vi_[m]; - } + Eigen::Map y_vi(y_vi_, M_[0], M_[1]); + y_vi = val_y.unaryExpr([](double x) { return new vari(x, false); }); + var_y.vi() = y_vi; return var_y; } @@ -405,14 +378,14 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template - inline void accumulate_adjoints(const Eigen::Matrix& y_adj_jac, + template ...> + inline void accumulate_adjoints(const T& y_adj_jac, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; if (is_var_[t]) { - for (int n = 0; n < y_adj_jac.size(); ++n) { - x_vis_[offsets_[t] + n]->adj_ += y_adj_jac(n); - } + Eigen::Map(&x_vis_[offsets_[t]], y_adj_jac.rows(), + y_adj_jac.cols()).adj() += y_adj_jac; } accumulate_adjoints(args...); } From 0edf1a2cdd9701c08e4f86c4e5d5bc5a19556b2a Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 31 Mar 2020 09:09:29 +0800 Subject: [PATCH 2/4] Generalise adj_jac_apply and softmax --- stan/math/fwd/fun/softmax.hpp | 58 ++++++++++----------- stan/math/prim/fun/softmax.hpp | 15 +++--- stan/math/rev/fun/softmax.hpp | 16 +++--- stan/math/rev/functor/adj_jac_apply.hpp | 69 +++++++++++++++++-------- test/unit/math/mix/fun/softmax_test.cpp | 63 ++++++++++++++++++++++ 5 files changed, 154 insertions(+), 67 deletions(-) diff --git a/stan/math/fwd/fun/softmax.hpp b/stan/math/fwd/fun/softmax.hpp index e30165f3cac..3330182603f 100644 --- a/stan/math/fwd/fun/softmax.hpp +++ b/stan/math/fwd/fun/softmax.hpp @@ -8,43 +8,37 @@ namespace stan { namespace math { -template -inline Eigen::Matrix, Eigen::Dynamic, 1> softmax( - const Eigen::Matrix, Eigen::Dynamic, 1>& alpha) { - using Eigen::Dynamic; - using Eigen::Matrix; - - Matrix alpha_t(alpha.size()); - for (int k = 0; k < alpha.size(); ++k) { - alpha_t(k) = alpha(k).val_; - } - - Matrix softmax_alpha_t = softmax(alpha_t); - - Matrix, Dynamic, 1> softmax_alpha(alpha.size()); - for (int k = 0; k < alpha.size(); ++k) { - softmax_alpha(k).val_ = softmax_alpha_t(k); - softmax_alpha(k).d_ = 0; - } - - for (int m = 0; m < alpha.size(); ++m) { - T negative_alpha_m_d_times_softmax_alpha_t_m - = -alpha(m).d_ * softmax_alpha_t(m); - for (int k = 0; k < alpha.size(); ++k) { - if (m == k) { - softmax_alpha(k).d_ - += softmax_alpha_t(k) - * (alpha(m).d_ + negative_alpha_m_d_times_softmax_alpha_t_m); - } else { - softmax_alpha(k).d_ - += negative_alpha_m_d_times_softmax_alpha_t_m * softmax_alpha_t(k); +template ...> +inline auto softmax(const Container& x) { + return apply_vector_unary::apply(x, [&](const auto& alpha) { + using T_fvar_inner = typename value_type_t::Scalar; + plain_type_t softmax_alpha(alpha.size()); + + softmax_alpha.val() = softmax(alpha.val()); + softmax_alpha.d().setZero(); + + for (int m = 0; m < alpha.size(); ++m) { + T_fvar_inner negative_alpha_m_d_times_softmax_alpha_t_m + = -alpha(m).d_ * softmax_alpha.val().coeff(m); + for (int k = 0; k < alpha.size(); ++k) { + if (m == k) { + softmax_alpha(k).d_ + += softmax_alpha.val().coeff(k) + * (alpha(m).d_ + negative_alpha_m_d_times_softmax_alpha_t_m); + } else { + softmax_alpha(k).d_ + += negative_alpha_m_d_times_softmax_alpha_t_m + * softmax_alpha.val().coeff(k); + } } } - } - return softmax_alpha; + return softmax_alpha; + }); } + } // namespace math } // namespace stan #endif diff --git a/stan/math/prim/fun/softmax.hpp b/stan/math/prim/fun/softmax.hpp index bfc6b8658fa..6305677b759 100644 --- a/stan/math/prim/fun/softmax.hpp +++ b/stan/math/prim/fun/softmax.hpp @@ -2,6 +2,7 @@ #define STAN_MATH_PRIM_FUN_SOFTMAX_HPP #include +#include #include #include @@ -9,7 +10,8 @@ namespace stan { namespace math { /** - * Return the softmax of the specified vector. + * Return the softmax of the specified Eigen expression/object, + * std::vector, or container of these. * *

* \f$ @@ -37,16 +39,17 @@ namespace math { * \end{array} * \f$ * - * @tparam T type of elements in the vector - * @param[in] v Vector to transform. - * @return Unit simplex result of the softmax transform of the vector. + * @tparam Container type of container + * @param[in] x Container (or container of containers) to transform. + * @return Unit simplex result of the softmax transform of the container. */ -template >...> +template ...> inline auto softmax(const Container& x) { return apply_vector_unary::apply(x, [](const auto& v) { const Eigen::Ref>& v_ref = v; check_nonzero_size("softmax", "v", v_ref); - plain_type_t theta(v_ref.size()); + plain_type_t theta(v_ref.size()); theta = (v_ref.array() - v_ref.maxCoeff()).exp(); return (theta.array() / theta.sum()).eval(); }); diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp index 3096d574d62..3aaf962351f 100644 --- a/stan/math/rev/fun/softmax.hpp +++ b/stan/math/rev/fun/softmax.hpp @@ -62,20 +62,22 @@ class softmax_op { } // namespace internal /** - * Return the softmax of the specified Eigen vector. Softmax is + * Return the softmax of the specified Eigen expression/object, + * std::vector, or container of these. Softmax is * guaranteed to return a simplex. * * @param alpha Unconstrained input vector. * @return Softmax of the input. * @throw std::domain_error If the input vector is size 0. */ -template >>...> -inline auto softmax(const T& x) { - return apply_vector_unary::apply(x, [&](const auto& alpha) { - const Eigen::Ref>& a_ref = alpha; - check_nonzero_size("softmax", "alpha", alpha); +template ...> +inline auto softmax(const Container& alpha) { + return apply_vector_unary::apply(alpha, [&](const auto& v) { + const Eigen::Ref>& v_ref = v; + check_nonzero_size("softmax", "alpha", v_ref); - return adj_jac_apply(a_ref); + return adj_jac_apply(v_ref); }); } diff --git a/stan/math/rev/functor/adj_jac_apply.hpp b/stan/math/rev/functor/adj_jac_apply.hpp index 8a39b769798..e9288600db1 100644 --- a/stan/math/rev/functor/adj_jac_apply.hpp +++ b/stan/math/rev/functor/adj_jac_apply.hpp @@ -48,6 +48,7 @@ inline void build_y_adj(vari** y_vi, const std::array& M, /** * Store the adjoints from y_vi in y_adj * + * @tparam EigT type of Eigen input * @tparam size dimensionality of M * @tparam R number of rows, can be Eigen::Dynamic * @tparam C number of columns, can be Eigen::Dynamic @@ -56,10 +57,11 @@ inline void build_y_adj(vari** y_vi, const std::array& M, * @param[in] M shape of y_adj * @param[out] y_adj reference to Eigen::Matrix where adjoints are to be stored */ -template ...> +template ...> inline void build_y_adj(vari** y_vi, const std::array& M, EigT& y_adj) { - y_adj = Eigen::Map(y_vi, M[0], M[1]).adj(); + y_adj.matrix() = Eigen::Map(y_vi, M[0], M[1]).adj(); } /** @@ -71,8 +73,10 @@ template struct compute_dims {}; /** - * Compute the dimensionality of the given template argument. Double + * Compute the dimensionality of the given template argument. Arithmetic * types have dimensionality zero. + * + * @tparam T arithmetic type of input */ template struct compute_dims> { @@ -82,6 +86,8 @@ struct compute_dims> { /** * Compute the dimensionality of the given template argument. * std::vector has dimension 1 + * + * @tparam T type of scalar within std::vector */ template struct compute_dims> { @@ -90,14 +96,12 @@ struct compute_dims> { /** * Compute the dimensionality of the given template argument. - * Eigen::Matrix types all have dimension two. + * Eigen types all have dimension two. * - * @tparam T type of elements in the matrix - * @tparam R number of rows, can be Eigen::Dynamic - * @tparam C number of columns, can be Eigen::Dynamic + * @tparam EigT type of Eigen input */ -template -struct compute_dims> { +template +struct compute_dims> { static constexpr size_t value = 2; }; } // namespace internal @@ -118,7 +122,8 @@ template struct adj_jac_vari : public vari { std::array is_var_; using FReturnType - = std::result_of_t()))...)>; + = std::result_of_t()))...)>; F f_; std::array offsets_; @@ -137,12 +142,10 @@ struct adj_jac_vari : public vari { * The array offsets_ is populated with values to indicate where in x_vis_ the * vari pointers for each argument will be stored. * - * Each of the arguments can be an Eigen::Matrix with var or double scalar + * Each of the arguments can be an Eigen type with var or double scalar * types, a std::vector with var, double, or int scalar types, or a var, a * double, or an int. * - * @tparam R number of rows, can be Eigen::Dynamic - * @tparam C number of columns, can be Eigen::Dynamic * @tparam Pargs Types of rest of arguments * * @param count rolling count of number of varis that must be allocated @@ -210,12 +213,10 @@ struct adj_jac_vari : public vari { * the index starting at offsets_[n]. For Eigen::Matrix types, this copying is * done in with column major ordering. * - * Each of the arguments can be an Eigen::Matrix with var or double scalar + * Each of the arguments can be an Eigen type with var or double scalar * types, a std::vector with var, double, or int scalar types, or a var, a * double, or an int. * - * @tparam R number of rows, can be Eigen::Dynamic - * @tparam C number of columns, can be Eigen::Dynamic * @tparam Pargs Types of the rest of the arguments to be processed * * @param x next argument to have its vari pointers copied if necessary @@ -230,7 +231,8 @@ struct adj_jac_vari : public vari { prepare_x_vis(args...); } - template ...> + template ...> inline void prepare_x_vis(const T& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -313,13 +315,11 @@ struct adj_jac_vari : public vari { * initialized with the values of val_y. The shape of the new matrix comes * from M_ * - * @tparam R number of rows, can be Eigen::Dynamic - * @tparam C number of columns, can be Eigen::Dynamic * @param val_y output of F::operator() * @return Eigen::Matrix of vars */ template , - require_t>...> + require_t>>...> inline auto build_return_varis_and_vars(const T& val_y) { constexpr int R = T_plain::RowsAtCompileTime; constexpr int C = T_plain::ColsAtCompileTime; @@ -336,6 +336,33 @@ struct adj_jac_vari : public vari { return var_y; } + /** + * Return an Eigen::Array of vars created from newly allocated varis + * initialized with the values of val_y. The shape of the new array comes + * from M_ + * + * @param val_y output of F::operator() + * @return Eigen::Matrix of vars + */ + template , + require_t>>...> + inline auto build_return_varis_and_vars(const T& val_y) { + constexpr int R = T_plain::RowsAtCompileTime; + constexpr int C = T_plain::ColsAtCompileTime; + M_[0] = val_y.rows(); + M_[1] = val_y.cols(); + Eigen::Array var_y(M_[0], M_[1]); + + y_vi_ + = ChainableStack::instance_->memalloc_.alloc_array(var_y.size()); + Eigen::Map y_vi(y_vi_, M_[0], M_[1]); + y_vi.array() = val_y + .unaryExpr([](double x) { return new vari(x, false); }); + var_y.vi().array() = y_vi.array(); + + return var_y; + } + void prepare_x_vis() {} /** @@ -370,8 +397,6 @@ struct adj_jac_vari : public vari { * of x_vis_. Recursively calls accumulate_adjoints on the rest of the * arguments. * - * @tparam R number of rows, can be Eigen::Dynamic - * @tparam C number of columns, can be Eigen::Dynamic * @tparam Pargs Types of the rest of adjoints to accumulate * * @param y_adj_jac set of values to be accumulated in adjoints diff --git a/test/unit/math/mix/fun/softmax_test.cpp b/test/unit/math/mix/fun/softmax_test.cpp index 9fc4ca29f3b..ec81113d695 100644 --- a/test/unit/math/mix/fun/softmax_test.cpp +++ b/test/unit/math/mix/fun/softmax_test.cpp @@ -1,4 +1,5 @@ #include +#include TEST(MathMixMatFun, softmax) { auto f = [](const auto& x) { return stan::math::softmax(x); }; @@ -33,4 +34,66 @@ TEST(MathMixMatFun, softmax) { Eigen::VectorXd d4(3); d4 << 0, 3, -1; stan::test::expect_ad(tols, f, d4); + + // Row Vectors + Eigen::RowVectorXd rx0(0); // error case + stan::test::expect_ad(tols, f, rx0); + + Eigen::RowVectorXd rx1(1); + rx1 << 0; + stan::test::expect_ad(tols, f, rx1); + + Eigen::RowVectorXd rx2(2); + rx2 << -1, 1; + stan::test::expect_ad(tols, f, rx2); + + Eigen::RowVectorXd rx3(3); + rx3 << -1, 1, 10; + stan::test::expect_ad(tols, f, rx3); + + Eigen::RowVectorXd rx3b(3); + rx3b << 0, 1, 2; + stan::test::expect_ad(tols, f, rx3b); + + Eigen::RowVectorXd rx3c(3); + rx3c << 2, 1, 1; + stan::test::expect_ad(tols, f, rx3c); + + // std vectors + std::vector stx0(0); // error case + stan::test::expect_ad(tols, f, stx0); + + std::vector stx1{0}; + stan::test::expect_ad(tols, f, stx1); + + std::vector stx2{-1, 1}; + stan::test::expect_ad(tols, f, stx2); + + std::vector stx3{-1, 1, 10}; + stan::test::expect_ad(tols, f, stx3); + + std::vector stx3b{0, 1, 2}; + stan::test::expect_ad(tols, f, stx3b); + + std::vector stx3c{2, 1, 1}; + stan::test::expect_ad(tols, f, stx3c); + + // Nested containers + std::vector stvx0{a, a}; // error case + stan::test::expect_ad(tols, f, stvx0); + + std::vector stvx1{b, b}; + stan::test::expect_ad(tols, f, stvx1); + + std::vector strx0{rx0, rx0}; // error case + stan::test::expect_ad(tols, f, strx0); + + std::vector strx1{rx1, rx1}; + stan::test::expect_ad(tols, f, strx1); + + std::vector> ststx0{stx0, stx0}; // error case + stan::test::expect_ad(tols, f, ststx0); + + std::vector> ststx1{stx1, stx1}; + stan::test::expect_ad(tols, f, ststx1); } From d979b544e09d5cea0f0823760fb4b9c5f85f9be2 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 31 Mar 2020 10:16:51 +0800 Subject: [PATCH 3/4] Update doc and tidying --- stan/math/fwd/fun/softmax.hpp | 15 ++- stan/math/prim/fun/softmax.hpp | 7 +- stan/math/rev/fun/softmax.hpp | 3 +- stan/math/rev/functor/adj_jac_apply.hpp | 124 ++++++++++-------------- 4 files changed, 65 insertions(+), 84 deletions(-) diff --git a/stan/math/fwd/fun/softmax.hpp b/stan/math/fwd/fun/softmax.hpp index 3330182603f..915437290a8 100644 --- a/stan/math/fwd/fun/softmax.hpp +++ b/stan/math/fwd/fun/softmax.hpp @@ -7,9 +7,17 @@ namespace stan { namespace math { - -template ...> +/** + * Return the softmax of the specified Eigen expression/object, + * std::vector, or container of these. Softmax is + * guaranteed to return a simplex. + * + * @tparam Container type of container. + * @param x Unconstrained input vector. + * @return Softmax of the input. + * @throw std::domain_error If the input vector is size 0. + */ +template ...> inline auto softmax(const Container& x) { return apply_vector_unary::apply(x, [&](const auto& alpha) { using T_fvar_inner = typename value_type_t::Scalar; @@ -38,7 +46,6 @@ inline auto softmax(const Container& x) { }); } - } // namespace math } // namespace stan #endif diff --git a/stan/math/prim/fun/softmax.hpp b/stan/math/prim/fun/softmax.hpp index 6305677b759..c1b24422c33 100644 --- a/stan/math/prim/fun/softmax.hpp +++ b/stan/math/prim/fun/softmax.hpp @@ -47,10 +47,11 @@ template ...> inline auto softmax(const Container& x) { return apply_vector_unary::apply(x, [](const auto& v) { - const Eigen::Ref>& v_ref = v; + using T_plain = plain_type_t; + const Eigen::Ref& v_ref = v; check_nonzero_size("softmax", "v", v_ref); - plain_type_t theta(v_ref.size()); - theta = (v_ref.array() - v_ref.maxCoeff()).exp(); + + T_plain theta = (v_ref.array() - v_ref.maxCoeff()).exp(); return (theta.array() / theta.sum()).eval(); }); } diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp index 3aaf962351f..d67fa056121 100644 --- a/stan/math/rev/fun/softmax.hpp +++ b/stan/math/rev/fun/softmax.hpp @@ -70,8 +70,7 @@ class softmax_op { * @return Softmax of the input. * @throw std::domain_error If the input vector is size 0. */ -template ...> +template ...> inline auto softmax(const Container& alpha) { return apply_vector_unary::apply(alpha, [&](const auto& v) { const Eigen::Ref>& v_ref = v; diff --git a/stan/math/rev/functor/adj_jac_apply.hpp b/stan/math/rev/functor/adj_jac_apply.hpp index e9288600db1..42f62703c8e 100644 --- a/stan/math/rev/functor/adj_jac_apply.hpp +++ b/stan/math/rev/functor/adj_jac_apply.hpp @@ -50,8 +50,6 @@ inline void build_y_adj(vari** y_vi, const std::array& M, * * @tparam EigT type of Eigen input * @tparam size dimensionality of M - * @tparam R number of rows, can be Eigen::Dynamic - * @tparam C number of columns, can be Eigen::Dynamic * * @param[in] y_vi pointer to pointers to varis * @param[in] M shape of y_adj @@ -76,10 +74,10 @@ struct compute_dims {}; * Compute the dimensionality of the given template argument. Arithmetic * types have dimensionality zero. * - * @tparam T arithmetic type of input + * @tparam ArithT arithmetic type of input */ -template -struct compute_dims> { +template +struct compute_dims> { static constexpr size_t value = 0; }; @@ -87,10 +85,10 @@ struct compute_dims> { * Compute the dimensionality of the given template argument. * std::vector has dimension 1 * - * @tparam T type of scalar within std::vector + * @tparam StdVecT type of scalar within std::vector */ -template -struct compute_dims> { +template +struct compute_dims> { static constexpr size_t value = 1; }; @@ -107,7 +105,7 @@ struct compute_dims> { } // namespace internal /** - * adj_jac_vari interfaces a user supplied functor with the reverse mode + * adj_jac_vari interfaces a user supplied functor with the reverse mode * autodiff. It allows someone to implement functions with custom reverse mode * autodiff without having to deal with autodiff types. * @@ -153,8 +151,8 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template ...> - inline size_t count_memory(size_t count, const T& x, + template ...> + inline size_t count_memory(size_t count, const EigT& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -162,9 +160,9 @@ struct adj_jac_vari : public vari { return count_memory(count, args...); } - template ...> - inline size_t count_memory(size_t count, const T& x, + template ...> + inline size_t count_memory(size_t count, const EigT& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -180,9 +178,9 @@ struct adj_jac_vari : public vari { return count_memory(count, args...); } - template ...> - inline size_t count_memory(size_t count, const T& x, + template ...> + inline size_t count_memory(size_t count, const StdVecT& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -197,8 +195,8 @@ struct adj_jac_vari : public vari { return count_memory(count, args...); } - template ...> - inline size_t count_memory(size_t count, const T& x, + template ...> + inline size_t count_memory(size_t count, const ArithT& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; offsets_[t] = count; @@ -223,17 +221,17 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template ...> - inline void prepare_x_vis(const T& x, const Pargs&... args) { + template ...> + inline void prepare_x_vis(const EigT& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; Eigen::Map(&x_vis_[offsets_[t]], x.rows(), x.cols()) = x.matrix().vi(); prepare_x_vis(args...); } - template ...> - inline void prepare_x_vis(const T& x, const Pargs&... args) { + template ...> + inline void prepare_x_vis(const EigT& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -246,9 +244,9 @@ struct adj_jac_vari : public vari { prepare_x_vis(args...); } - template ...> - inline void prepare_x_vis(const T& x, + template ...> + inline void prepare_x_vis(const StdVecT& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -260,8 +258,8 @@ struct adj_jac_vari : public vari { prepare_x_vis(args...); } - template ...> - inline void prepare_x_vis(const T& x, const Pargs&... args) { + template ...> + inline void prepare_x_vis(const ArithT& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -315,14 +313,15 @@ struct adj_jac_vari : public vari { * initialized with the values of val_y. The shape of the new matrix comes * from M_ * + * @tparam EigT Type of Eigen input * @param val_y output of F::operator() * @return Eigen::Matrix of vars */ - template , - require_t>>...> - inline auto build_return_varis_and_vars(const T& val_y) { - constexpr int R = T_plain::RowsAtCompileTime; - constexpr int C = T_plain::ColsAtCompileTime; + template , + require_t>...> + inline auto build_return_varis_and_vars(const EigT& val_y) { + constexpr int R = EigT_plain::RowsAtCompileTime; + constexpr int C = EigT_plain::ColsAtCompileTime; M_[0] = val_y.rows(); M_[1] = val_y.cols(); Eigen::Matrix var_y(M_[0], M_[1]); @@ -341,14 +340,15 @@ struct adj_jac_vari : public vari { * initialized with the values of val_y. The shape of the new array comes * from M_ * + * @tparam EigT Type of Eigen input * @param val_y output of F::operator() * @return Eigen::Matrix of vars */ - template , - require_t>>...> - inline auto build_return_varis_and_vars(const T& val_y) { - constexpr int R = T_plain::RowsAtCompileTime; - constexpr int C = T_plain::ColsAtCompileTime; + template , + require_t>...> + inline auto build_return_varis_and_vars(const EigT& val_y) { + constexpr int R = EigT_plain::RowsAtCompileTime; + constexpr int C = EigT_plain::ColsAtCompileTime; M_[0] = val_y.rows(); M_[1] = val_y.cols(); Eigen::Array var_y(M_[0], M_[1]); @@ -397,15 +397,16 @@ struct adj_jac_vari : public vari { * of x_vis_. Recursively calls accumulate_adjoints on the rest of the * arguments. * + * @tparam EigT Type of Eigen input * @tparam Pargs Types of the rest of adjoints to accumulate * * @param y_adj_jac set of values to be accumulated in adjoints * @param args the rest of the arguments (that will be iterated through * recursively) */ - template ...> - inline void accumulate_adjoints(const T& y_adj_jac, + template ...> + inline void accumulate_adjoints(const EigT& y_adj_jac, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; if (is_var_[t]) { @@ -426,8 +427,9 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template - inline void accumulate_adjoints(const std::vector& y_adj_jac, + template ...> + inline void accumulate_adjoints(const StdVecT& y_adj_jac, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; if (is_var_[t]) { @@ -439,21 +441,6 @@ struct adj_jac_vari : public vari { accumulate_adjoints(args...); } - /** - * Recursively call accumulate_adjoints with args. There are no adjoints to - * accumulate for std::vector arguments. - * - * @tparam Pargs Types of the rest of adjoints to accumulate - * @param y_adj_jac ignored - * @param args the rest of the arguments (that will be iterated through - * recursively) - */ - template - inline void accumulate_adjoints(const std::vector& y_adj_jac, - const Pargs&... args) { - accumulate_adjoints(args...); - } - /** * Accumulate, if necessary, the value of y_adj_jac into the * adjoint of the vari pointed to by the appropriate element @@ -465,8 +452,9 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template - inline void accumulate_adjoints(const double& y_adj_jac, + template ...> + inline void accumulate_adjoints(const ArithT& y_adj_jac, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; if (is_var_[t]) { @@ -475,20 +463,6 @@ struct adj_jac_vari : public vari { accumulate_adjoints(args...); } - /** - * Recursively call accumulate_adjoints with args. There are no adjoints to - * accumulate for int arguments. - * - * @tparam Pargs Types of the rest of adjoints to accumulate - * @param y_adj_jac ignored - * @param args the rest of the arguments (that will be iterated through - * recursively) - */ - template - inline void accumulate_adjoints(const int& y_adj_jac, const Pargs&... args) { - accumulate_adjoints(args...); - } - inline void accumulate_adjoints() {} /** From b809db4ffb8021cc3c695950e54ebe081eabde13 Mon Sep 17 00:00:00 2001 From: Stan Jenkins Date: Mon, 30 Mar 2020 22:30:17 -0400 Subject: [PATCH 4/4] [Jenkins] auto-formatting by clang-format version 6.0.0 (tags/google/stable/2017-11-14) --- stan/math/fwd/fun/softmax.hpp | 5 ++--- stan/math/opencl/kernel_generator/load.hpp | 2 +- stan/math/prim/fun/value_of.hpp | 2 +- stan/math/rev/fun/softmax.hpp | 5 ++--- stan/math/rev/functor/adj_jac_apply.hpp | 25 ++++++++++------------ 5 files changed, 17 insertions(+), 22 deletions(-) diff --git a/stan/math/fwd/fun/softmax.hpp b/stan/math/fwd/fun/softmax.hpp index 915437290a8..472b264b143 100644 --- a/stan/math/fwd/fun/softmax.hpp +++ b/stan/math/fwd/fun/softmax.hpp @@ -35,9 +35,8 @@ inline auto softmax(const Container& x) { += softmax_alpha.val().coeff(k) * (alpha(m).d_ + negative_alpha_m_d_times_softmax_alpha_t_m); } else { - softmax_alpha(k).d_ - += negative_alpha_m_d_times_softmax_alpha_t_m - * softmax_alpha.val().coeff(k); + softmax_alpha(k).d_ += negative_alpha_m_d_times_softmax_alpha_t_m + * softmax_alpha.val().coeff(k); } } } diff --git a/stan/math/opencl/kernel_generator/load.hpp b/stan/math/opencl/kernel_generator/load.hpp index 799cdb551a0..ae0f7d1eb6d 100644 --- a/stan/math/opencl/kernel_generator/load.hpp +++ b/stan/math/opencl/kernel_generator/load.hpp @@ -49,7 +49,7 @@ class load_ * Creates a deep copy of this expression. * @return copy of \c *this */ - inline load_ deep_copy() const & { return load_(a_); } + inline load_ deep_copy() const& { return load_(a_); } inline load_ deep_copy() && { return load_(std::forward(a_)); } /** diff --git a/stan/math/prim/fun/value_of.hpp b/stan/math/prim/fun/value_of.hpp index a86d0d02915..2641cbdca5a 100644 --- a/stan/math/prim/fun/value_of.hpp +++ b/stan/math/prim/fun/value_of.hpp @@ -120,7 +120,7 @@ inline const std::vector& value_of(const std::vector& x) { return x; } **/ template ...> inline decltype(auto) value_of(const T& M) { - return M.unaryExpr([](const auto& x){ return value_of(x); }).eval(); + return M.unaryExpr([](const auto& x) { return value_of(x); }).eval(); } /** diff --git a/stan/math/rev/fun/softmax.hpp b/stan/math/rev/fun/softmax.hpp index d67fa056121..8bf06dfa9f0 100644 --- a/stan/math/rev/fun/softmax.hpp +++ b/stan/math/rev/fun/softmax.hpp @@ -28,7 +28,7 @@ class softmax_op { */ template plain_type_t operator()(const std::array& /* needs_adj */, - const T& alpha) { + const T& alpha) { N_ = alpha.size(); y_ = ChainableStack::instance_->memalloc_.alloc_array(N_); Eigen::Map> y_map(y_, N_); @@ -49,8 +49,7 @@ class softmax_op { */ template > std::tuple multiply_adjoint_jacobian( - const std::array& /* needs_adj */, - const T& adj) const { + const std::array& /* needs_adj */, const T& adj) const { T_plain adj_times_jac(N_); Eigen::Map y(y_, N_); diff --git a/stan/math/rev/functor/adj_jac_apply.hpp b/stan/math/rev/functor/adj_jac_apply.hpp index 42f62703c8e..7804068a0ad 100644 --- a/stan/math/rev/functor/adj_jac_apply.hpp +++ b/stan/math/rev/functor/adj_jac_apply.hpp @@ -119,9 +119,8 @@ struct compute_dims> { template struct adj_jac_vari : public vari { std::array is_var_; - using FReturnType - = std::result_of_t()))...)>; + using FReturnType = std::result_of_t()))...)>; F f_; std::array offsets_; @@ -225,7 +224,7 @@ struct adj_jac_vari : public vari { inline void prepare_x_vis(const EigT& x, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; Eigen::Map(&x_vis_[offsets_[t]], x.rows(), x.cols()) - = x.matrix().vi(); + = x.matrix().vi(); prepare_x_vis(args...); } @@ -246,8 +245,7 @@ struct adj_jac_vari : public vari { template ...> - inline void prepare_x_vis(const StdVecT& x, - const Pargs&... args) { + inline void prepare_x_vis(const StdVecT& x, const Pargs&... args) { prepare_x_vis(args...); } @@ -329,7 +327,7 @@ struct adj_jac_vari : public vari { y_vi_ = ChainableStack::instance_->memalloc_.alloc_array(var_y.size()); Eigen::Map y_vi(y_vi_, M_[0], M_[1]); - y_vi = val_y.unaryExpr([](double x) { return new vari(x, false); }); + y_vi = val_y.unaryExpr([](double x) { return new vari(x, false); }); var_y.vi() = y_vi; return var_y; @@ -356,8 +354,7 @@ struct adj_jac_vari : public vari { y_vi_ = ChainableStack::instance_->memalloc_.alloc_array(var_y.size()); Eigen::Map y_vi(y_vi_, M_[0], M_[1]); - y_vi.array() = val_y - .unaryExpr([](double x) { return new vari(x, false); }); + y_vi.array() = val_y.unaryExpr([](double x) { return new vari(x, false); }); var_y.vi().array() = y_vi.array(); return var_y; @@ -406,12 +403,13 @@ struct adj_jac_vari : public vari { */ template ...> - inline void accumulate_adjoints(const EigT& y_adj_jac, - const Pargs&... args) { + inline void accumulate_adjoints(const EigT& y_adj_jac, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1; if (is_var_[t]) { Eigen::Map(&x_vis_[offsets_[t]], y_adj_jac.rows(), - y_adj_jac.cols()).adj() += y_adj_jac; + y_adj_jac.cols()) + .adj() + += y_adj_jac; } accumulate_adjoints(args...); } @@ -452,8 +450,7 @@ struct adj_jac_vari : public vari { * @param args the rest of the arguments (that will be iterated through * recursively) */ - template ...> + template ...> inline void accumulate_adjoints(const ArithT& y_adj_jac, const Pargs&... args) { static constexpr int t = sizeof...(Targs) - sizeof...(Pargs) - 1;