Unify weighted mean code

cpp/include/raft/stats/detail/weighted_mean.cuh

@@ -18,14 +18,17 @@
 
 #include <raft/cudart_utils.h>
 #include <raft/linalg/coalesced_reduction.hpp>
+#include <raft/linalg/reduce.cuh>
 #include <raft/linalg/strided_reduction.hpp>
+#include <raft/stats/sum.cuh>
 
 namespace raft {
 namespace stats {
 namespace detail {
 
 /**
- * @brief Compute the row-wise weighted mean of the input matrix
+ * @brief Compute the row-wise weighted mean of the input matrix with a
+ * vector of column weights
  *
  * @tparam Type the data type
  * @param mu the output mean vector
@@ -58,7 +61,52 @@ void rowWeightedMean(
 }
 
 /**
- * @brief Compute the column-wise weighted mean of the input matrix
+ * @brief Compute the row-wise weighted mean of the input matrix with a
+ * vector of sample weights
+ *
+ * @tparam Type the data type
+ * @param mu the output mean vector
+ * @param data the input matrix
+ * @param weights per-sample weight
+ * @param D number of columns of data
+ * @param N number of rows of data
+ * @param row_major input matrix is row-major or not
+ * @param along_rows whether to reduce along rows or columns
+ * @param stream cuda stream to launch work on
+ */
+template <typename Type>
+void rowSampleWeightedMean(Type* mu,
+                           const Type* data,
+                           const Type* weights,
+                           int D,
+                           int N,
+                           bool row_major,
+                           bool along_rows,
+                           cudaStream_t stream)
+{
+  // sum the weights & copy back to CPU
+  Type WS = 0;
+  raft::stats::sum(mu, weights, 1, N, row_major, stream);
+  raft::update_host(&WS, mu, 1, stream);
+
+  raft::linalg::reduce(
+    mu,
+    data,
+    D,
+    N,
+    (Type)0,
+    row_major,
+    along_rows,
+    stream,
+    false,
+    [weights] __device__(Type v, int i) { return v * weights[i]; },
+    [] __device__(Type a, Type b) { return a + b; },
+    [WS] __device__(Type v) { return v / WS; });
+}
+
+/**
+ * @brief Compute the column-wise weighted mean of the input matrix with a
+ * vector of column weights
  *
  * @tparam Type the data type
  * @param mu the output mean vector

cpp/include/raft/stats/weighted_mean.hpp

@@ -22,7 +22,8 @@ namespace raft {
 namespace stats {
 
 /**
- * @brief Compute the row-wise weighted mean of the input matrix
+ * @brief Compute the row-wise weighted mean of the input matrix with a
+ * vector of column weights
  *
  * @tparam Type the data type
  * @param mu the output mean vector
@@ -40,7 +41,35 @@ void rowWeightedMean(
 }
 
 /**
- * @brief Compute the column-wise weighted mean of the input matrix
+ * @brief Compute the row-wise weighted mean of the input matrix with a
+ * vector of sample weights
+ *
+ * @tparam Type the data type
+ * @param mu the output mean vector
+ * @param data the input matrix
+ * @param weights per-sample weight
+ * @param D number of columns of data
+ * @param N number of rows of data
+ * @param row_major input matrix is row-major or not
+ * @param along_rows whether to reduce along rows or columns
+ * @param stream cuda stream to launch work on
+ */
+template <typename Type>
+void rowSampleWeightedMean(Type* mu,
+                           const Type* data,
+                           const Type* weights,
+                           int D,
+                           int N,
+                           bool row_major,
+                           bool along_rows,
+                           cudaStream_t stream)
+{
+  detail::rowSampleWeightedMean(mu, data, weights, D, N, row_major, along_rows, stream);
+}
+
+/**
+ * @brief Compute the column-wise weighted mean of the input matrix with a
+ * vector of column weights
  *
  * @tparam Type the data type
  * @param mu the output mean vector