Merge pull request #3376 from SteveBronder/test/reduce-sum-tuple-regressions-3041-3359

Allow `reduce_sum` to take in tuples and tuples of tuples

Author: WardBrian

stan/math/rev/core/accumulate_adjoints.hpp

@@ -1,6 +1,7 @@
 #ifndef STAN_MATH_REV_CORE_ACCUMULATE_ADJOINTS_HPP
 #define STAN_MATH_REV_CORE_ACCUMULATE_ADJOINTS_HPP
 
+#include <stan/math/prim/functor/apply.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core/var.hpp>
@@ -33,6 +34,9 @@ template <typename Arith, require_st_arithmetic<Arith>* = nullptr,
           typename... Pargs>
 inline double* accumulate_adjoints(double* dest, Arith&& x, Pargs&&... args);
 
+template <typename Tuple, require_tuple_t<Tuple>* = nullptr, typename... Pargs>
+inline double* accumulate_adjoints(double* dest, Tuple&& x, Pargs&&... args);
+
 inline double* accumulate_adjoints(double* dest);
 
 /**
@@ -140,6 +144,27 @@ inline double* accumulate_adjoints(double* dest, Arith&& x, Pargs&&... args) {
   return accumulate_adjoints(dest, std::forward<Pargs>(args)...);
 }
 
+/**
+ * Accumulate adjoints from a tuple into storage pointed to by dest, then
+ * recursively accumulate adjoints from the remaining arguments.
+ *
+ * @tparam Tuple A tuple type
+ * @tparam Pargs Types of remaining arguments
+ * @param dest Pointer to where adjoints are to be accumulated
+ * @param x A tuple containing arguments whose adjoints are accumulated
+ * @param args Further args to accumulate over
+ * @return Final position of adjoint storage pointer
+ */
+template <typename Tuple, require_tuple_t<Tuple>*, typename... Pargs>
+inline double* accumulate_adjoints(double* dest, Tuple&& x, Pargs&&... args) {
+  dest = stan::math::apply(
+      [dest](auto&&... tuple_args) {
+        return accumulate_adjoints(dest, tuple_args...);
+      },
+      std::forward<Tuple>(x));
+  return accumulate_adjoints(dest, std::forward<Pargs>(args)...);
+}
+
 /**
  * End accumulate_adjoints recursion and return pointer
  *

stan/math/rev/core/count_vars.hpp

@@ -33,6 +33,9 @@ template <typename Arith, require_arithmetic_t<scalar_type_t<Arith>>* = nullptr,
           typename... Pargs>
 inline size_t count_vars_impl(size_t count, Arith& x, Pargs&&... args);
 
+template <typename Tuple, require_tuple_t<Tuple>* = nullptr, typename... Pargs>
+inline size_t count_vars_impl(size_t count, Tuple&& x, Pargs&&... args);
+
 inline size_t count_vars_impl(size_t count);
 /**
  * Count the number of vars in x (a std::vector of vars),
@@ -132,23 +135,33 @@ inline size_t count_vars_impl(size_t count, Arith& x, Pargs&&... args) {
 inline size_t count_vars_impl(size_t count, std::basic_ostream<char>*&) {
   return count;
 }
+
+/**
+ * Count the vars in a tuple, add them to the running total, and count the vars
+ * in the remaining arguments.
+ *
+ * @tparam Tuple A tuple type
+ * @tparam Pargs Types of remaining arguments
+ * @param[in] count The current count of the number of vars
+ * @param[in] x A tuple containing arguments to count
+ * @param[in] args Objects to be forwarded to the recursive call
+ * @return The total number of vars
+ */
+template <typename Tuple, require_tuple_t<Tuple>*, typename... Pargs>
+inline size_t count_vars_impl(size_t count, Tuple&& x, Pargs&&... args) {
+  count = stan::math::apply(
+      [count](auto&&... tuple_args) {
+        return count_vars_impl(count, tuple_args...);
+      },
+      std::forward<Tuple>(x));
+  return count_vars_impl(count, std::forward<Pargs>(args)...);
+}
+
 /**
  * End count_vars_impl recursion and return total number of counted vars
  */
 inline size_t count_vars_impl(size_t count) { return count; }
 
-template <typename... Pargs, typename... Args>
-inline size_t count_vars_impl(std::size_t count,
-                              const std::tuple<Pargs...>& arg, Args&&... args) {
-  return count_vars_impl(
-      stan::math::apply(
-          [count](auto&&... inner_args) {
-            return (count_vars_impl(0, inner_args) + ... + count);
-          },
-          arg),
-      std::forward<Args>(args)...);
-}
-
 }  // namespace internal
 
 /**

stan/math/rev/core/deep_copy_vars.hpp

@@ -1,6 +1,8 @@
 #ifndef STAN_MATH_REV_CORE_DEEP_COPY_VARS_HPP
 #define STAN_MATH_REV_CORE_DEEP_COPY_VARS_HPP
 
+#include <stan/math/prim/functor/apply.hpp>
+#include <stan/math/prim/functor/make_holder_tuple.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core/var.hpp>
@@ -81,6 +83,23 @@ inline auto deep_copy_vars(EigT&& arg) {
       .eval();
 }
 
+/**
+ * Copy the vars in a tuple but reallocate new varis for them.
+ *
+ * @tparam Tuple A tuple type
+ * @param arg A tuple containing arguments to copy
+ * @return A tuple containing copied arguments
+ */
+template <typename Tuple, require_tuple_t<Tuple>* = nullptr>
+inline auto deep_copy_vars(Tuple&& arg) {
+  return stan::math::apply(
+      [](auto&&... tuple_args) {
+        return make_holder_tuple(
+            deep_copy_vars(std::forward<decltype(tuple_args)>(tuple_args))...);
+      },
+      std::forward<Tuple>(arg));
+}
+
 }  // namespace math
 }  // namespace stan
 

stan/math/rev/core/save_varis.hpp

@@ -2,6 +2,7 @@
 #define STAN_MATH_REV_CORE_SAVE_VARIS_HPP
 
 #include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/functor/apply.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/rev/meta.hpp>
 #include <stan/math/rev/core/var.hpp>
@@ -33,6 +34,9 @@ template <typename Arith, require_st_arithmetic<Arith>* = nullptr,
           typename... Pargs>
 inline vari** save_varis(vari** dest, Arith&& x, Pargs&&... args);
 
+template <typename Tuple, require_tuple_t<Tuple>* = nullptr, typename... Pargs>
+inline vari** save_varis(vari** dest, Tuple&& x, Pargs&&... args);
+
 inline vari** save_varis(vari** dest);
 
 /**
@@ -136,6 +140,25 @@ inline vari** save_varis(vari** dest, Arith&& x, Pargs&&... args) {
   return save_varis(dest, std::forward<Pargs>(args)...);
 }
 
+/**
+ * Save the vari pointers in a tuple into the memory pointed to by dest, then
+ * recursively save the varis in the remaining arguments.
+ *
+ * @tparam Tuple A tuple type
+ * @tparam Pargs Types of remaining arguments
+ * @param[in, out] dest Pointer to where vari pointers are saved
+ * @param[in] x A tuple containing arguments whose varis are saved
+ * @param[in] args Additional arguments to have their varis saved
+ * @return Final position of dest pointer
+ */
+template <typename Tuple, require_tuple_t<Tuple>*, typename... Pargs>
+inline vari** save_varis(vari** dest, Tuple&& x, Pargs&&... args) {
+  dest = stan::math::apply(
+      [dest](auto&&... tuple_args) { return save_varis(dest, tuple_args...); },
+      std::forward<Tuple>(x));
+  return save_varis(dest, std::forward<Pargs>(args)...);
+}
+
 /**
  * End save_varis recursion and return pointer
  *

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

@@ -2,6 +2,7 @@
 #include <stan/math/rev/core.hpp>
 #include <stan/math.hpp>
 #include <test/unit/math/rev/util.hpp>
+#include <tuple>
 #include <vector>
 
 TEST_F(AgradRev, Rev_accumulate_adjoints_zero_args) {
@@ -437,3 +438,39 @@ TEST_F(AgradRev, Rev_accumulate_adjoints_sum) {
   EXPECT_EQ(ptr, storage.data() + num_vars);
   stan::math::recover_memory();
 }
+
+TEST_F(AgradRev, Rev_accumulate_adjoints_tuple_args) {
+  const std::tuple<> empty;
+  const auto data = std::make_tuple(1, Eigen::VectorXd::Ones(2));
+  Eigen::VectorXd data_storage = Eigen::VectorXd::Zero(2);
+  double* data_ptr
+      = stan::math::accumulate_adjoints(data_storage.data(), empty, data);
+  EXPECT_EQ(data_storage.data(), data_ptr);
+  EXPECT_FLOAT_EQ(0.0, data_storage(0));
+  EXPECT_FLOAT_EQ(0.0, data_storage(1));
+
+  stan::math::var before = 1.0;
+  stan::math::var first = 2.0;
+  Eigen::Matrix<stan::math::var, Eigen::Dynamic, 1> vars(2);
+  vars << 3.0, 4.0;
+  stan::math::var last = 5.0;
+  stan::math::var after = 6.0;
+  before.vi_->adj_ = 1.0;
+  first.vi_->adj_ = 2.0;
+  vars(0).vi_->adj_ = 3.0;
+  vars(1).vi_->adj_ = 4.0;
+  last.vi_->adj_ = 5.0;
+  after.vi_->adj_ = 6.0;
+  auto nested = std::make_tuple(first, std::make_tuple(vars, 7), last);
+  Eigen::VectorXd storage = Eigen::VectorXd::Constant(8, 10.0);
+
+  double* ptr = stan::math::accumulate_adjoints(storage.data(), before, nested,
+                                                std::make_tuple(after));
+
+  EXPECT_EQ(storage.data() + 6, ptr);
+  for (int i = 0; i < 6; ++i) {
+    EXPECT_FLOAT_EQ(11.0 + i, storage(i));
+  }
+  EXPECT_FLOAT_EQ(10.0, storage(6));
+  EXPECT_FLOAT_EQ(10.0, storage(7));
+}

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

@@ -3,6 +3,7 @@
 #include <stan/math.hpp>
 #include <test/unit/math/rev/util.hpp>
 
+#include <tuple>
 #include <vector>
 
 using stan::math::var;
@@ -161,3 +162,18 @@ TEST_F(AgradRev, Rev_count_vars_sum) {
       count_vars(arg1, arg18, arg17, arg2, arg16, arg3, arg15, arg4, arg14,
                  arg5, arg13, arg12, arg6, arg11, arg7, arg10, arg8, arg9));
 }
+
+TEST_F(AgradRev, Rev_count_vars_tuple_args) {
+  const std::tuple<> empty;
+  const auto data = std::make_tuple(1, Eigen::VectorXd::Ones(2));
+  EXPECT_EQ(0, stan::math::count_vars(empty));
+  EXPECT_EQ(0, stan::math::count_vars(data));
+
+  Eigen::Matrix<var, Eigen::Dynamic, 1> vars(2);
+  auto nested = std::make_tuple(var(1.0), std::make_tuple(vars, 2.0), var(3.0));
+  var before = 4.0;
+
+  EXPECT_EQ(5, stan::math::count_vars(before, nested));
+  EXPECT_EQ(2, stan::math::count_vars(
+                   std::make_tuple(var(5.0), std::make_tuple(var(6.0)))));
+}

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

@@ -2,6 +2,8 @@
 #include <stan/math/rev/core.hpp>
 #include <stan/math.hpp>
 #include <test/unit/math/rev/util.hpp>
+#include <tuple>
+#include <type_traits>
 #include <vector>
 
 using stan::math::var;
@@ -306,3 +308,45 @@ TEST_F(AgradRev, Rev_deep_copy_vars_std_vector_eigen_matrix_var_arg) {
       EXPECT_NE(out[i](j).vi_, arg[i](j).vi_);
     }
 }
+
+TEST_F(AgradRev, Rev_deep_copy_vars_tuple_data_arg) {
+  const std::tuple<> empty;
+  const auto arg = std::make_tuple(5, Eigen::VectorXd::Ones(2).eval());
+
+  auto empty_out = stan::math::deep_copy_vars(empty);
+  auto out = stan::math::deep_copy_vars(arg);
+
+  static_assert(std::is_same_v<decltype(empty_out), std::tuple<>>);
+  static_assert(std::is_same_v<decltype(out),
+                               std::tuple<const int&, const Eigen::VectorXd&>>);
+  EXPECT_EQ(&std::get<0>(out), &std::get<0>(arg));
+  EXPECT_EQ(&std::get<1>(out), &std::get<1>(arg));
+}
+
+TEST_F(AgradRev, Rev_deep_copy_vars_nested_tuple_var_arg) {
+  Eigen::Matrix<var, Eigen::Dynamic, 1> vars(2);
+  vars << 2.0, 3.0;
+  auto arg = std::make_tuple(1.0, var(4.0), vars, std::make_tuple(var(5.0), 6));
+
+  auto out = stan::math::deep_copy_vars(arg);
+
+  static_assert(std::is_reference_v<std::tuple_element_t<0, decltype(out)>>);
+  EXPECT_EQ(&std::get<0>(out), &std::get<0>(arg));
+  EXPECT_FLOAT_EQ(std::get<1>(out).val(), std::get<1>(arg).val());
+  EXPECT_NE(std::get<1>(out).vi_, std::get<1>(arg).vi_);
+  for (int i = 0; i < vars.size(); ++i) {
+    EXPECT_FLOAT_EQ(std::get<2>(out)(i).val(), std::get<2>(arg)(i).val());
+    EXPECT_NE(std::get<2>(out)(i).vi_, std::get<2>(arg)(i).vi_);
+  }
+  EXPECT_FLOAT_EQ(std::get<0>(std::get<3>(out)).val(),
+                  std::get<0>(std::get<3>(arg)).val());
+  EXPECT_NE(std::get<0>(std::get<3>(out)).vi_,
+            std::get<0>(std::get<3>(arg)).vi_);
+  EXPECT_EQ(&std::get<1>(std::get<3>(out)), &std::get<1>(std::get<3>(arg)));
+
+  auto rvalue_out = stan::math::deep_copy_vars(
+      std::make_tuple(7.0, var(8.0), std::make_tuple(9)));
+  static_assert(std::is_same_v<decltype(rvalue_out),
+                               std::tuple<double, var, std::tuple<int>>>);
+  EXPECT_FLOAT_EQ(8.0, std::get<1>(rvalue_out).val());
+}

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

@@ -1,6 +1,7 @@
 #include <stan/math.hpp>
 #include <test/unit/math/rev/util.hpp>
 #include <gtest/gtest.h>
+#include <tuple>
 #include <vector>
 
 using stan::math::var;
@@ -424,3 +425,34 @@ TEST_F(AgradRev, Rev_save_varis_sum) {
 
   EXPECT_EQ(ptr, storage.data() + num_vars);
 }
+
+TEST_F(AgradRev, Rev_save_varis_tuple_args) {
+  const std::tuple<> empty;
+  const auto data = std::make_tuple(1, Eigen::VectorXd::Ones(2));
+  std::vector<vari*> data_storage(2, nullptr);
+  vari** data_ptr = stan::math::save_varis(data_storage.data(), empty, data);
+  EXPECT_EQ(data_storage.data(), data_ptr);
+  EXPECT_EQ(nullptr, data_storage[0]);
+  EXPECT_EQ(nullptr, data_storage[1]);
+
+  var before = 1.0;
+  var first = 2.0;
+  Eigen::Matrix<var, Eigen::Dynamic, 1> vars(2);
+  vars << 3.0, 4.0;
+  var last = 5.0;
+  var after = 6.0;
+  auto nested = std::make_tuple(first, std::make_tuple(vars, 7), last);
+  std::vector<vari*> storage(8, nullptr);
+
+  vari** ptr = stan::math::save_varis(storage.data(), before, nested,
+                                      std::make_tuple(after));
+
+  std::vector<vari*> expected{before.vi_,  first.vi_, vars(0).vi_,
+                              vars(1).vi_, last.vi_,  after.vi_};
+  EXPECT_EQ(storage.data() + expected.size(), ptr);
+  for (size_t i = 0; i < expected.size(); ++i) {
+    EXPECT_EQ(expected[i], storage[i]);
+  }
+  EXPECT_EQ(nullptr, storage[expected.size()]);
+  EXPECT_EQ(nullptr, storage[expected.size() + 1]);
+}

test/unit/math/rev/functor/reduce_sum_tuple_data_test.cpp

@@ -0,0 +1,34 @@
+#include <stan/math.hpp>
+#include <test/unit/math/rev/util.hpp>
+#include <gtest/gtest.h>
+#include <tuple>
+#include <vector>
+
+namespace {
+
+struct sum_with_laplace_options {
+  template <typename Options>
+  auto operator()(const std::vector<int>& slice, std::size_t start,
+                  std::size_t end, std::ostream* msgs,
+                  const stan::math::var& shared,
+                  const Options& laplace_options) const {
+    return slice.size()
+           * (shared + static_cast<double>(std::get<2>(laplace_options)));
+  }
+};
+
+// Regression test for https://github.com/stan-dev/math/issues/3359.
+TEST_F(AgradRev, reduce_sum_accepts_const_laplace_options_tuple) {
+  stan::math::var shared = 1.0;
+  const auto laplace_options = stan::math::generate_laplace_options(1);
+
+  stan::math::var result = stan::math::reduce_sum<sum_with_laplace_options>(
+      std::vector<int>{0, 1}, 1, nullptr, shared, laplace_options);
+
+  const double expected = 2.0 * (shared.val() + std::get<2>(laplace_options));
+  EXPECT_FLOAT_EQ(expected, result.val());
+  result.grad();
+  EXPECT_FLOAT_EQ(2.0, shared.adj());
+}
+
+}  // namespace