Stan Math Library  2.20.0
reverse mode automatic differentiation
softmax.hpp
Go to the documentation of this file.
1 #ifndef STAN_MATH_REV_MAT_FUN_SOFTMAX_HPP
2 #define STAN_MATH_REV_MAT_FUN_SOFTMAX_HPP
3 
4 #include <stan/math/rev/meta.hpp>
8 #include <vector>
9 #include <tuple>
10 
11 namespace stan {
12 namespace math {
13 
14 namespace internal {
15 class softmax_op {
16  int N_;
17  double* y_; // Holds the results of the softmax
18 
19  public:
20  softmax_op() : N_(0), y_(NULL) {}
21 
22  /*
23  * Compute the softmax of the unconstrained input vector
24  *
25  * @param alpha Unconstrained input vector.
26  * @return Softmax of the input.
27  */
28  template <std::size_t size>
29  Eigen::VectorXd operator()(const std::array<bool, size>& needs_adj,
30  const Eigen::VectorXd& alpha) {
31  N_ = alpha.size();
33 
34  auto y = softmax(alpha);
35  for (int n = 0; n < N_; ++n)
36  y_[n] = y(n);
37  return y;
38  }
39 
40  /*
41  * Compute the result of multiply the transpose of the adjoint vector times
42  * the Jacobian of the softmax operator. It is more efficient to do this
43  * without actually computing the Jacobian and doing the vector-matrix
44  * product.
45  *
46  * @param adj Eigen::VectorXd of adjoints at the output of the softmax
47  * @return Eigen::VectorXd of adjoints propagated through softmax operation
48  */
49  template <std::size_t size>
50  std::tuple<Eigen::VectorXd> multiply_adjoint_jacobian(
51  const std::array<bool, size>& needs_adj,
52  const Eigen::VectorXd& adj) const {
53  Eigen::VectorXd adj_times_jac(N_);
54  Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, 1> > y(y_, N_);
55 
56  double adj_dot_y = adj.dot(y);
57  for (int n = 0; n < N_; ++n) {
58  adj_times_jac(n) = -y(n) * adj_dot_y + y(n) * adj(n);
59  }
60 
61  return std::make_tuple(adj_times_jac);
62  }
63 };
64 } // namespace internal
65 
74 inline Eigen::Matrix<var, Eigen::Dynamic, 1> softmax(
75  const Eigen::Matrix<var, Eigen::Dynamic, 1>& alpha) {
76  check_nonzero_size("softmax", "alpha", alpha);
77 
78  return adj_jac_apply<internal::softmax_op>(alpha);
79 }
80 
81 } // namespace math
82 } // namespace stan
83 #endif
void check_nonzero_size(const char *function, const char *name, const T_y &y)
Check if the specified matrix/vector is of non-zero size.
Eigen::Matrix< fvar< T >, Eigen::Dynamic, 1 > softmax(const Eigen::Matrix< fvar< T >, Eigen::Dynamic, 1 > &alpha)
Definition: softmax.hpp:12
Eigen::VectorXd operator()(const std::array< bool, size > &needs_adj, const Eigen::VectorXd &alpha)
Definition: softmax.hpp:29
static STAN_THREADS_DEF AutodiffStackStorage * instance_
std::tuple< Eigen::VectorXd > multiply_adjoint_jacobian(const std::array< bool, size > &needs_adj, const Eigen::VectorXd &adj) const
Definition: softmax.hpp:50
T * alloc_array(size_t n)
Allocate an array on the arena of the specified size to hold values of the specified template paramet...

     [ Stan Home Page ] © 2011–2018, Stan Development Team.