Merge pull request #2230 from stan-dev/feature/varmat-mdivide-left

Added `var<mat>` implementation of `mdivide_left` (Issue #2101)

Author: bbbales2

stan/math/rev/core.hpp

@@ -7,6 +7,7 @@
 #include <stan/math/rev/core/autodiffstackstorage.hpp>
 #include <stan/math/rev/core/build_vari_array.hpp>
 #include <stan/math/rev/core/chainable_alloc.hpp>
+#include <stan/math/rev/core/chainable_object.hpp>
 #include <stan/math/rev/core/chainablestack.hpp>
 #include <stan/math/rev/core/count_vars.hpp>
 #include <stan/math/rev/core/callback_vari.hpp>

stan/math/rev/core/chainable_object.hpp

@@ -0,0 +1,67 @@
+#ifndef STAN_MATH_REV_CORE_CHAINABLE_OBJECT_HPP
+#define STAN_MATH_REV_CORE_CHAINABLE_OBJECT_HPP
+
+#include <stan/math/rev/meta.hpp>
+#include <stan/math/rev/core/chainable_alloc.hpp>
+#include <stan/math/rev/core/typedefs.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/typedefs.hpp>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/**
+ * `chainable_object` hold another object is useful for connecting
+ * the lifetime of a specific object to the chainable stack
+ *
+ * `chainable_object` objects should only be allocated with `new`.
+ * `chainable_object` objects allocated on the stack will result
+ * in a double free (`obj_` will get destructed once when the
+ * chainable_object leaves scope and once when the chainable
+ * stack memory is recovered).
+ *
+ * @tparam T type of object to hold
+ */
+template <typename T>
+class chainable_object : public chainable_alloc {
+ private:
+  plain_type_t<T> obj_;
+
+ public:
+  /**
+   * Construct chainable object from another object
+   *
+   * @tparam S type of object to hold (must have the same plain type as `T`)
+   */
+  template <typename S,
+            require_same_t<plain_type_t<T>, plain_type_t<S>>* = nullptr>
+  explicit chainable_object(S&& obj) : obj_(std::forward<S>(obj)) {}
+
+  /**
+   * Return a reference to the underlying object
+   *
+   * @return reference to underlying object
+   */
+  inline auto& get() noexcept { return obj_; }
+  inline const auto& get() const noexcept { return obj_; }
+};
+
+/**
+ * Store the given object in a `chainable_object` so it is destructed
+ * only when the chainable stack memory is recovered and return
+ * a pointer to the underlying object
+ *
+ * @tparam T type of object to hold
+ * @param obj object to hold
+ * @return pointer to object held in `chainable_object`
+ */
+template <typename T>
+auto make_chainable_ptr(T&& obj) {
+  auto ptr = new chainable_object<T>(std::forward<T>(obj));
+  return &ptr->get();
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/rev/fun/mdivide_left.hpp

@@ -1,10 +1,10 @@
-
 #ifndef STAN_MATH_REV_FUN_MDIVIDE_LEFT_HPP
 #define STAN_MATH_REV_FUN_MDIVIDE_LEFT_HPP
 
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core.hpp>
 #include <stan/math/rev/core/typedefs.hpp>
+#include <stan/math/rev/core/chainable_object.hpp>
 #include <stan/math/prim/err.hpp>
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/typedefs.hpp>
@@ -13,45 +13,74 @@
 namespace stan {
 namespace math {
 
-template <typename T1, typename T2, require_all_eigen_t<T1, T2>* = nullptr,
-          require_any_vt_var<T1, T2>* = nullptr>
+/**
+ * Return the solution `X` of `AX = B`.
+ *
+ * A must be a square matrix, but B can be a matrix or a vector
+ *
+ * @tparam T1 type of first matrix
+ * @tparam T2 type of second matrix
+ *
+ * @param[in] A square matrix
+ * @param[in] B right hand side
+ * @return solution of AX = B
+ */
+template <typename T1, typename T2, require_all_matrix_t<T1, T2>* = nullptr,
+          require_any_st_var<T1, T2>* = nullptr>
 inline auto mdivide_left(const T1& A, const T2& B) {
-  using ret_type = plain_type_t<decltype(A * B)>;
+  using ret_val_type = plain_type_t<decltype(value_of(A) * value_of(B))>;
+  using ret_type = promote_var_matrix_t<ret_val_type, T1, T2>;
 
   check_square("mdivide_left", "A", A);
   check_multiplicable("mdivide_left", "A", A, "B", B);
 
   if (A.size() == 0) {
-    return ret_type(0, B.cols());
+    return ret_type(ret_val_type(0, B.cols()));
   }
+
   if (!is_constant<T1>::value && !is_constant<T2>::value) {
     arena_t<promote_scalar_t<var, T1>> arena_A = A;
     arena_t<promote_scalar_t<var, T2>> arena_B = B;
-    arena_t<promote_scalar_t<double, T1>> arena_A_val = arena_A.val();
-    arena_t<ret_type> res = arena_A_val.householderQr().solve(arena_B.val());
-    reverse_pass_callback([arena_A, arena_B, arena_A_val, res]() mutable {
+
+    auto hqr_A_ptr = make_chainable_ptr(arena_A.val().householderQr());
+    arena_t<ret_type> res = hqr_A_ptr->solve(arena_B.val());
+    reverse_pass_callback([arena_A, arena_B, hqr_A_ptr, res]() mutable {
       promote_scalar_t<double, T2> adjB
-          = arena_A_val.transpose().householderQr().solve(res.adj());
+          = hqr_A_ptr->householderQ()
+            * hqr_A_ptr->matrixQR()
+                  .template triangularView<Eigen::Upper>()
+                  .transpose()
+                  .solve(res.adj());
       arena_A.adj() -= adjB * res.val_op().transpose();
       arena_B.adj() += adjB;
     });
 
     return ret_type(res);
   } else if (!is_constant<T2>::value) {
     arena_t<promote_scalar_t<var, T2>> arena_B = B;
-    arena_t<promote_scalar_t<double, T1>> arena_A_val = value_of(A);
-    arena_t<ret_type> res = arena_A_val.householderQr().solve(arena_B.val());
-    reverse_pass_callback([arena_B, arena_A_val, res]() mutable {
-      arena_B.adj() += arena_A_val.transpose().householderQr().solve(res.adj());
+
+    auto hqr_A_ptr = make_chainable_ptr(value_of(A).householderQr());
+    arena_t<ret_type> res = hqr_A_ptr->solve(arena_B.val());
+    reverse_pass_callback([arena_B, hqr_A_ptr, res]() mutable {
+      arena_B.adj() += hqr_A_ptr->householderQ()
+                       * hqr_A_ptr->matrixQR()
+                             .template triangularView<Eigen::Upper>()
+                             .transpose()
+                             .solve(res.adj());
     });
     return ret_type(res);
   } else {
     arena_t<promote_scalar_t<var, T1>> arena_A = A;
-    arena_t<ret_type> res = arena_A.val().householderQr().solve(value_of(B));
-    reverse_pass_callback([arena_A, res]() mutable {
-      arena_A.adj()
-          -= arena_A.val().transpose().householderQr().solve(res.adj())
-             * res.val_op().transpose();
+
+    auto hqr_A_ptr = make_chainable_ptr(arena_A.val().householderQr());
+    arena_t<ret_type> res = hqr_A_ptr->solve(value_of(B));
+    reverse_pass_callback([arena_A, hqr_A_ptr, res]() mutable {
+      arena_A.adj() -= hqr_A_ptr->householderQ()
+                       * hqr_A_ptr->matrixQR()
+                             .template triangularView<Eigen::Upper>()
+                             .transpose()
+                             .solve(res.adj())
+                       * res.val_op().transpose();
     });
     return ret_type(res);
   }

test/unit/math/mix/fun/mdivide_left_test.cpp

@@ -13,18 +13,24 @@ TEST(MathMixMatFun, mdivideLeft) {
   stan::test::expect_ad(f, m00, m00);
   stan::test::expect_ad(f, m00, m02);
   stan::test::expect_ad(f, m00, v0);
+  stan::test::expect_ad_matvar(f, m00, m00);
+  stan::test::expect_ad_matvar(f, m00, m02);
+  stan::test::expect_ad_matvar(f, m00, v0);
 
   Eigen::MatrixXd aa(1, 1);
   aa << 1;
   Eigen::MatrixXd bb(1, 1);
   bb << 2;
   stan::test::expect_ad(f, aa, bb);
+  stan::test::expect_ad_matvar(f, aa, bb);
   Eigen::MatrixXd b0(1, 0);
   stan::test::expect_ad(f, aa, b0);
+  stan::test::expect_ad_matvar(f, aa, b0);
 
   Eigen::VectorXd cc(1);
   cc << 3;
   stan::test::expect_ad(f, aa, cc);
+  stan::test::expect_ad_matvar(f, aa, cc);
 
   Eigen::MatrixXd a(2, 2);
   a << 2, 3, 3, 7;
@@ -47,12 +53,14 @@ TEST(MathMixMatFun, mdivideLeft) {
   for (const auto& m1 : std::vector<Eigen::MatrixXd>{a, b, c, d}) {
     for (const auto& m2 : std::vector<Eigen::MatrixXd>{a, b, c, d, e}) {
       stan::test::expect_ad(f, m1, m2);
+      stan::test::expect_ad_matvar(f, m1, m2);
     }
   }
 
   // matrix, vector
   for (const auto& m : std::vector<Eigen::MatrixXd>{a, b, c, d}) {
     stan::test::expect_ad(f, m, g);
+    stan::test::expect_ad_matvar(f, m, g);
   }
 
   Eigen::MatrixXd v(5, 5);
@@ -61,6 +69,7 @@ TEST(MathMixMatFun, mdivideLeft) {
   Eigen::VectorXd u(5);
   u << 62, 84, 84, 76, 108;
   stan::test::expect_ad(f, v, u);
+  stan::test::expect_ad_matvar(f, v, u);
 
   Eigen::MatrixXd m33 = Eigen::MatrixXd::Zero(3, 3);
   Eigen::MatrixXd m44 = Eigen::MatrixXd::Zero(4, 4);
@@ -72,7 +81,10 @@ TEST(MathMixMatFun, mdivideLeft) {
   // exceptions: wrong sizes
   stan::test::expect_ad(f, m33, m44);
   stan::test::expect_ad(f, m33, v4);
+  stan::test::expect_ad_matvar(f, m33, m44);
+  stan::test::expect_ad_matvar(f, m33, v4);
 
   // exceptions: wrong types
   stan::test::expect_ad(f, m33, rv3);
+  stan::test::expect_ad_matvar(f, m33, rv3);
 }

test/unit/math/mix/prob/multi_gp_cholesky_test.cpp

@@ -44,6 +44,8 @@ TEST(ProbDistributionsMultiGPCholesky, fvar_var) {
                   stan::math::multi_gp_cholesky_log(y, L, w).val_.val());
   EXPECT_FLOAT_EQ(-74.572952,
                   stan::math::multi_gp_cholesky_log(y, L, w).d_.val());
+
+  stan::math::recover_memory();
 }
 
 TEST(ProbDistributionsMultiGPCholesky, fvar_fvar_var) {
@@ -89,4 +91,6 @@ TEST(ProbDistributionsMultiGPCholesky, fvar_fvar_var) {
                   stan::math::multi_gp_cholesky_log(y, L, w).val_.val_.val());
   EXPECT_FLOAT_EQ(-74.572952,
                   stan::math::multi_gp_cholesky_log(y, L, w).d_.val_.val());
+
+  stan::math::recover_memory();
 }

test/unit/math/rev/core/chainable_object_test.cpp

@@ -0,0 +1,68 @@
+#include <stan/math.hpp>
+#include <gtest/gtest.h>
+#include <test/unit/util.hpp>
+
+class ChainableObjectTest {
+ public:
+  static int counter;
+
+  ~ChainableObjectTest() { counter++; }
+};
+
+int ChainableObjectTest::counter = 0;
+
+TEST(AgradRev, chainable_object_test) {
+  {
+    auto ptr = new stan::math::chainable_object<ChainableObjectTest>(
+        ChainableObjectTest());
+    ChainableObjectTest::counter = 0;
+  }
+
+  EXPECT_EQ((ChainableObjectTest::counter), 0);
+  stan::math::recover_memory();
+  EXPECT_EQ((ChainableObjectTest::counter), 1);
+}
+
+TEST(AgradRev, chainable_object_nested_test) {
+  stan::math::start_nested();
+
+  {
+    auto ptr = new stan::math::chainable_object<ChainableObjectTest>(
+        ChainableObjectTest());
+    ChainableObjectTest::counter = 0;
+  }
+
+  EXPECT_EQ((ChainableObjectTest::counter), 0);
+
+  stan::math::recover_memory_nested();
+
+  EXPECT_EQ((ChainableObjectTest::counter), 1);
+}
+
+TEST(AgradRev, make_chainable_ptr_test) {
+  {
+    ChainableObjectTest* ptr
+        = stan::math::make_chainable_ptr(ChainableObjectTest());
+    ChainableObjectTest::counter = 0;
+  }
+
+  EXPECT_EQ((ChainableObjectTest::counter), 0);
+  stan::math::recover_memory();
+  EXPECT_EQ((ChainableObjectTest::counter), 1);
+}
+
+TEST(AgradRev, make_chainable_ptr_nested_test) {
+  stan::math::start_nested();
+
+  {
+    ChainableObjectTest* ptr
+        = stan::math::make_chainable_ptr(ChainableObjectTest());
+    ChainableObjectTest::counter = 0;
+  }
+
+  EXPECT_EQ((ChainableObjectTest::counter), 0);
+
+  stan::math::recover_memory_nested();
+
+  EXPECT_EQ((ChainableObjectTest::counter), 1);
+}