Automatic Differentiation
 
Loading...
Searching...
No Matches
diag_pre_multiply.hpp
Go to the documentation of this file.
1#ifndef STAN_MATH_REV_FUN_DIAG_PRE_MULTIPLY_HPP
2#define STAN_MATH_REV_FUN_DIAG_PRE_MULTIPLY_HPP
3
4#include <stan/math/prim/fun/Eigen.hpp>
5#include <stan/math/rev/core.hpp>
6#include <stan/math/prim/err.hpp>
7#include <stan/math/prim/fun/diag_pre_multiply.hpp>
8
9namespace stan {
10namespace math {
11
24template <typename T1, typename T2, require_vector_t<T1>* = nullptr,
25 require_matrix_t<T2>* = nullptr,
26 require_any_st_var<T1, T2>* = nullptr>
27inline auto diag_pre_multiply(const T1& m1, const T2& m2) {
28 check_size_match("diag_pre_multiply", "m1.size()", m1.size(), "m2.rows()",
29 m2.rows());
30 using inner_ret_type = decltype(value_of(m1).asDiagonal() * value_of(m2));
31 using ret_type = return_var_matrix_t<inner_ret_type, T1, T2>;
32 if constexpr (is_autodiff_v<T1> && is_autodiff_v<T2>) {
33 arena_t<promote_scalar_t<var, T1>> arena_m1 = m1;
34 arena_t<promote_scalar_t<var, T2>> arena_m2 = m2;
35 arena_t<ret_type> ret(arena_m1.val().asDiagonal() * arena_m2.val());
36 reverse_pass_callback([ret, arena_m1, arena_m2]() mutable {
37 arena_m1.adj() += arena_m2.val().cwiseProduct(ret.adj()).rowwise().sum();
38 arena_m2.adj() += arena_m1.val().asDiagonal() * ret.adj();
39 });
40 return ret_type(ret);
41 } else if constexpr (is_autodiff_v<T1>) {
42 arena_t<promote_scalar_t<var, T1>> arena_m1 = m1;
43 arena_t<promote_scalar_t<double, T2>> arena_m2 = value_of(m2);
44 arena_t<ret_type> ret(arena_m1.val().asDiagonal() * arena_m2);
45 reverse_pass_callback([ret, arena_m1, arena_m2]() mutable {
46 arena_m1.adj() += arena_m2.val().cwiseProduct(ret.adj()).rowwise().sum();
47 });
48 return ret_type(ret);
49 } else if constexpr (is_autodiff_v<T2>) {
50 arena_t<promote_scalar_t<double, T1>> arena_m1 = value_of(m1);
51 arena_t<promote_scalar_t<var, T2>> arena_m2 = m2;
52 arena_t<ret_type> ret(arena_m1.asDiagonal() * arena_m2.val());
53 reverse_pass_callback([ret, arena_m1, arena_m2]() mutable {
54 arena_m2.adj() += arena_m1.val().asDiagonal() * ret.adj();
55 });
56 return ret_type(ret);
57 }
58}
59
60} // namespace math
61} // namespace stan
62
63#endif
stan::math::diag_pre_multiply
auto diag_pre_multiply(const T1 &m1, const T2 &m2)
Return the product of the diagonal matrix formed from the vector or row_vector and a matrix.
Definition diag_pre_multiply.hpp:25
stan::math::reverse_pass_callback
void reverse_pass_callback(F &&functor)
Puts a callback on the autodiff stack to be called in reverse pass.
Definition reverse_pass_callback.hpp:38
stan::math::value_of
T value_of(const fvar< T > &v)
Return the value of the specified variable.
Definition value_of.hpp:18
stan::math::check_size_match
void check_size_match(const char *function, const char *name_i, T_size1 i, const char *name_j, T_size2 j)
Check if the provided sizes match.
Definition check_size_match.hpp:24
stan::arena_t
typename internal::arena_type_impl< std::decay_t< T > >::type arena_t
Determines a type that can be used in place of T that does any dynamic allocations on the AD stack.
Definition arena_type.hpp:60
stan::return_var_matrix_t
std::conditional_t< is_any_var_matrix< ReturnType, Types... >::value, stan::math::var_value< stan::math::promote_scalar_t< double, plain_type_t< ReturnType > > >, stan::math::promote_scalar_t< stan::math::var_value< double >, plain_type_t< ReturnType > > > return_var_matrix_t
Given an Eigen type and several inputs, determine if a matrix should be var<Matrix> or Matrix<var>.
Definition return_var_matrix.hpp:27
stan
The lgamma implementation in stan-math is based on either the reentrant safe lgamma_r implementation ...
Definition unit_vector_constrain.hpp:15
diag_pre_multiply.hpp