Merge branch 'develop' into feature/issue-54-dirichlet-multinomial

.github/workflows/header_checks.yml

@@ -27,7 +27,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.x'
 

.github/workflows/main.yml

@@ -25,7 +25,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - uses: r-lib/actions/setup-r@v2
@@ -75,7 +75,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - uses: r-lib/actions/setup-r@v2
@@ -129,7 +129,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - uses: r-lib/actions/setup-r@v2
@@ -178,7 +178,7 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
-    - uses: actions/setup-python@v4
+    - uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - uses: r-lib/actions/setup-r@v2

make/libraries

@@ -169,7 +169,7 @@ endif
 	touch $(TBB_BIN)/tbb-make-check
 
 
-$(TBB_BIN)/tbb.def: $(TBB_BIN)/tbb-make-check $(TBB_BIN)/tbbmalloc.def
+$(TBB_BIN)/tbb.def: $(TBB_BIN)/tbb-make-check
 	@mkdir -p $(TBB_BIN)
 	touch $(TBB_BIN)/version_$(notdir $(TBB))
 	tbb_root="$(TBB_RELATIVE_PATH)" CXX="$(CXX)" CC="$(TBB_CC)" LDFLAGS='$(LDFLAGS_TBB)' '$(MAKE)' -C "$(TBB_BIN)" -r -f "$(TBB_ABSOLUTE_PATH)/build/Makefile.tbb" compiler=$(TBB_CXX_TYPE) cfg=release stdver=c++1y  CXXFLAGS="$(TBB_CXXFLAGS)"

stan/math/opencl/matrix_cl.hpp

@@ -501,6 +501,37 @@ class matrix_cl : public matrix_cl_base {
    */
   ~matrix_cl() { wait_for_read_write_events(); }
 
+  /**
+   * Set the values of a `matrix_cl` to zero.
+   */
+  void setZero() {
+    if (this->size() == 0) {
+      return;
+    }
+    cl::Event zero_event;
+    const std::size_t write_events_size = this->write_events().size();
+    const std::size_t read_events_size = this->read_events().size();
+    const std::size_t read_write_size = write_events_size + read_events_size;
+    std::vector<cl::Event> read_write_events(read_write_size, cl::Event{});
+    auto&& read_events_vec = this->read_events();
+    auto&& write_events_vec = this->write_events();
+    for (std::size_t i = 0; i < read_events_size; ++i) {
+      read_write_events[i] = read_events_vec[i];
+    }
+    for (std::size_t i = read_events_size, j = 0; j < write_events_size;
+         ++i, ++j) {
+      read_write_events[i] = write_events_vec[j];
+    }
+    try {
+      opencl_context.queue().enqueueFillBuffer(buffer_cl_, static_cast<T>(0), 0,
+                                               sizeof(T) * this->size(),
+                                               &read_write_events, &zero_event);
+    } catch (const cl::Error& e) {
+      check_opencl_error("setZero", e);
+    }
+    this->add_write_event(zero_event);
+  }
+
  private:
   /**
    * Initializes the OpenCL buffer of this matrix by copying the data from given

stan/math/prim/fun/chol2inv.hpp

@@ -5,7 +5,7 @@
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/fun/dot_self.hpp>
 #include <stan/math/prim/fun/dot_product.hpp>
-#include <stan/math/prim/fun/mdivide_left_tri_low.hpp>
+#include <stan/math/prim/fun/mdivide_left_tri.hpp>
 #include <stan/math/prim/fun/inv_square.hpp>
 
 namespace stan {
@@ -35,7 +35,7 @@ plain_type_t<T> chol2inv(const T& L) {
     X.coeffRef(0) = inv_square(L_ref.coeff(0, 0));
     return X;
   }
-  T_result L_inv = mdivide_left_tri_low(L_ref, T_result::Identity(K, K));
+  T_result L_inv = mdivide_left_tri<Eigen::Lower>(L_ref);
   T_result X(K, K);
   for (int k = 0; k < K; ++k) {
     X.coeffRef(k, k) = dot_self(L_inv.col(k).tail(K - k).eval());

stan/math/prim/prob/inv_wishart_cholesky_rng.hpp

@@ -3,9 +3,7 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/mdivide_left_tri.hpp>
-#include <stan/math/prim/prob/wishart_cholesky_rng.hpp>
-#include <stan/math/prim/prob/wishart_rng.hpp>
+#include <stan/math/prim/fun/mdivide_left_tri_low.hpp>
 
 namespace stan {
 namespace math {
@@ -16,6 +14,9 @@ namespace math {
  * from the inverse Wishart distribution with the specified degrees of freedom
  * using the specified random number generator.
  *
+ * Axen, Seth D. "Efficiently generating inverse-Wishart matrices and their
+ * Cholesky factors." arXiv preprint arXiv:2310.15884 (2023).
+ *
  * @tparam RNG Random number generator type
  * @param[in] nu scalar degrees of freedom
  * @param[in] L_S lower Cholesky factor of the scale matrix
@@ -38,8 +39,15 @@ inline Eigen::MatrixXd inv_wishart_cholesky_rng(double nu,
   check_positive(function, "Cholesky Scale matrix", L_S.diagonal());
   check_positive(function, "columns of Cholesky Scale matrix", L_S.cols());
 
-  MatrixXd L_Sinv = mdivide_left_tri<Eigen::Lower>(L_S);
-  return mdivide_left_tri<Eigen::Lower>(wishart_cholesky_rng(nu, L_Sinv, rng));
+  MatrixXd B = MatrixXd::Zero(k, k);
+  for (int j = 0; j < k; ++j) {
+    for (int i = 0; i < j; ++i) {
+      B(j, i) = normal_rng(0, 1, rng);
+    }
+    B(j, j) = std::sqrt(chi_square_rng(nu - k + j + 1, rng));
+  }
+
+  return mdivide_left_tri_low(B, L_S);
 }
 
 }  // namespace math

stan/math/rev/core/arena_matrix.hpp

@@ -128,6 +128,15 @@ class arena_matrix : public Eigen::Map<MatrixType> {
     Base::operator=(a);
     return *this;
   }
+  /**
+   * Forces hard copying matrices into an arena matrix
+   * @tparam T Any type assignable to `Base`
+   * @param x the values to write to `this`
+   */
+  template <typename T>
+  void deep_copy(const T& x) {
+    Base::operator=(x);
+  }
 };
 
 }  // namespace math

stan/math/rev/core/var.hpp

@@ -1020,9 +1020,10 @@ class var_value<T, internal::require_matrix_var_value<T>> {
    * @param other the value to assign
    * @return this
    */
-  template <typename S, require_assignable_t<value_type, S>* = nullptr,
-            require_all_plain_type_t<T, S>* = nullptr,
-            require_not_same_t<plain_type_t<T>, plain_type_t<S>>* = nullptr>
+  template <typename S, typename T_ = T,
+            require_assignable_t<value_type, S>* = nullptr,
+            require_all_plain_type_t<T_, S>* = nullptr,
+            require_not_same_t<plain_type_t<T_>, plain_type_t<S>>* = nullptr>
   inline var_value<T>& operator=(const var_value<S>& other) {
     static_assert(
         EIGEN_PREDICATE_SAME_MATRIX_SIZE(T, S),
@@ -1032,16 +1033,65 @@ class var_value<T, internal::require_matrix_var_value<T>> {
   }
 
   /**
-   * Assignment of another var value, when either this or the other one does not
+   * Assignment of another var value, when the `this` does not
    * contain a plain type.
-   * @tparam S type of the value in the `var_value` to assing
+   * @tparam S type of the value in the `var_value` to assign
+   * @param other the value to assign
+   * @return this
+   */
+  template <typename S, typename T_ = T,
+            require_assignable_t<value_type, S>* = nullptr,
+            require_not_plain_type_t<S>* = nullptr,
+            require_plain_type_t<T_>* = nullptr>
+  inline var_value<T>& operator=(const var_value<S>& other) {
+    // If vi_ is nullptr then the var needs initialized via copy constructor
+    if (!(this->vi_)) {
+      *this = var_value<T>(other);
+      return *this;
+    }
+    arena_t<plain_type_t<T>> prev_val(vi_->val_.rows(), vi_->val_.cols());
+    prev_val.deep_copy(vi_->val_);
+    vi_->val_.deep_copy(other.val());
+    // no need to change any adjoints - these are just zeros before the reverse
+    // pass
+
+    reverse_pass_callback(
+        [this_vi = this->vi_, other_vi = other.vi_, prev_val]() mutable {
+          this_vi->val_.deep_copy(prev_val);
+
+          // we have no way of detecting aliasing between this->vi_->adj_ and
+          // other.vi_->adj_, so we must copy adjoint before reseting to zero
+
+          // we can reuse prev_val instead of allocating a new matrix
+          prev_val.deep_copy(this_vi->adj_);
+          this_vi->adj_.setZero();
+          other_vi->adj_ += prev_val;
+        });
+    return *this;
+  }
+  /**
+   * Assignment of another var value, when either both `this` or other does not
+   * contain a plain type.
+   * @note Here we do not need to use `deep_copy` as the `var_value`'s
+   * inner `vari_type` holds a view which will call the assignment operator
+   *  that does not perform a placement new.
+   * @tparam S type of the value in the `var_value` to assign
    * @param other the value to assign
    * @return this
    */
   template <typename S, typename T_ = T,
             require_assignable_t<value_type, S>* = nullptr,
-            require_any_not_plain_type_t<T_, S>* = nullptr>
+            require_not_plain_type_t<T_>* = nullptr>
   inline var_value<T>& operator=(const var_value<S>& other) {
+    // If vi_ is nullptr then the var needs initialized via copy constructor
+    if (!(this->vi_)) {
+      []() STAN_COLD_PATH {
+        throw std::domain_error(
+            "var_value<matrix>::operator=(var_value<expression>):"
+            " Internal Bug! Please report this with an example"
+            " of your model to the Stan math github repository.");
+      }();
+    }
     arena_t<plain_type_t<T>> prev_val = vi_->val_;
     vi_->val_ = other.val();
     // no need to change any adjoints - these are just zeros before the reverse

test/unit/math/opencl/matrix_cl_test.cpp

@@ -77,4 +77,18 @@ TEST(MathMatrixCL, assignment) {
   EXPECT_EQ(nullptr, mat1_cl.buffer()());
 }
 
+TEST(MathMatrixCL, setZeroFun) {
+  using stan::math::matrix_cl;
+  Eigen::Matrix<double, 2, 2> mat_1;
+  mat_1 << 1, 2, 3, 4;
+  matrix_cl<double> mat1_cl(mat_1);
+  mat1_cl.setZero();
+  Eigen::Matrix<double, 2, 2> mat_1_fromcl
+      = stan::math::from_matrix_cl(mat1_cl);
+  EXPECT_EQ(mat_1_fromcl(0), 0);
+  EXPECT_EQ(mat_1_fromcl(1), 0);
+  EXPECT_EQ(mat_1_fromcl(2), 0);
+  EXPECT_EQ(mat_1_fromcl(3), 0);
+}
+
 #endif

test/unit/math/prim/prob/inv_wishart_cholesky_rng_test.cpp

@@ -91,14 +91,15 @@ TEST(ProbDistributionsInvWishartCholesky, SpecialRNGTest) {
   using stan::math::inv_wishart_cholesky_rng;
   using stan::math::multiply_lower_tri_self_transpose;
 
-  boost::random::mt19937 rng(1234U);
+  boost::random::mt19937 rng(92343U);
   int N = 1e5;
   double tol = 0.1;
   for (int k = 1; k < 5; k++) {
-    MatrixXd sigma = MatrixXd::Identity(k, k);
+    MatrixXd L = MatrixXd::Identity(k, k);
     MatrixXd Z = MatrixXd::Zero(k, k);
     for (int i = 0; i < N; i++) {
-      Z += stan::math::crossprod(inv_wishart_cholesky_rng(k + 2, sigma, rng));
+      Z += multiply_lower_tri_self_transpose(
+          inv_wishart_cholesky_rng(k + 2, L, rng));
     }
     Z /= N;
     for (int j = 0; j < k; j++) {
@@ -111,3 +112,35 @@ TEST(ProbDistributionsInvWishartCholesky, SpecialRNGTest) {
     }
   }
 }
+
+TEST(ProbDistributionsInvWishartCholesky, compareToInvWishart) {
+  // Compare the marginal mean
+
+  using Eigen::MatrixXd;
+  using Eigen::VectorXd;
+  using stan::math::inv_wishart_cholesky_rng;
+  using stan::math::inv_wishart_rng;
+  using stan::math::multiply_lower_tri_self_transpose;
+  using stan::math::qr_thin_Q;
+
+  boost::random::mt19937 rng(92343U);
+  int N = 1e5;
+  double tol = 0.05;
+  for (int k = 1; k < 5; k++) {
+    MatrixXd L = qr_thin_Q(MatrixXd::Random(k, k)).transpose();
+    L.diagonal() = stan::math::abs(L.diagonal());
+    MatrixXd sigma = multiply_lower_tri_self_transpose(L);
+    MatrixXd Z_mean = sigma / (k + 3);
+    MatrixXd Z_est = MatrixXd::Zero(k, k);
+    for (int i = 0; i < N; i++) {
+      Z_est += multiply_lower_tri_self_transpose(
+          inv_wishart_cholesky_rng(k + 4, L, rng));
+    }
+    Z_est /= N;
+    for (int j = 0; j < k; j++) {
+      for (int i = 0; i < j; i++) {
+        EXPECT_NEAR(Z_est(i, j), Z_mean(i, j), tol);
+      }
+    }
+  }
+}