transpose matmul categorical bit reproducibility

Author: adityagoel4512

src/tabmat/ext/cat_split_helpers-tmpl.cpp

@@ -1,5 +1,5 @@
 #include <vector>
-
+#include <omp.h>
 
 <%def name="transpose_matvec(dropfirst)">
 template <typename Int, typename F>
@@ -10,24 +10,30 @@ void _transpose_matvec_${dropfirst}(
     F* res,
     Int res_size
 ) {
-    #pragma omp parallel
+    int num_threads = omp_get_max_threads();
+    std::vector<F> all_res(num_threads * res_size, 0.0);
+    #pragma omp parallel shared(all_res) 
     {
-        std::vector<F> restemp(res_size, 0.0);
-        #pragma omp for
+	int tid = omp_get_thread_num();
+	F* res_slice = &all_res[tid * res_size]; 
+	#pragma omp for
         for (Py_ssize_t i = 0; i < n_rows; i++) {
             % if dropfirst == 'all_rows_drop_first':
                 Py_ssize_t col_idx = indices[i] - 1;
                 if (col_idx != -1) {
-                    restemp[col_idx] += other[i];
+                    res_slice[col_idx] += other[i];
                 }
             % else:
-                restemp[indices[i]] += other[i];
+  	        res_slice[indices[i]] += other[i];
             % endif
         }
-        for (Py_ssize_t i = 0; i < res_size; i++) {
-            # pragma omp atomic
-            res[i] += restemp[i];
-        }
+	#pragma omp for
+	for (Py_ssize_t i = 0; i < res_size; ++i) {
+	    for (int tid = 0; tid < num_threads; ++tid) {
+		#pragma omp atomic
+	        res[i] += all_res[tid * res_size + i];	    
+	    }
+	}	
     }
 }
 </%def>

tests/test_reproducibility.py

@@ -0,0 +1,24 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+import tabmat as tm
+
+N = 100
+K = 5
+
+
+@pytest.fixture
+def df():
+    rng = np.random.default_rng(1234)
+    return pd.DataFrame(
+        pd.Categorical(rng.integers(low=0, high=K - 1, size=N), categories=range(K))
+    )
+
+
+@pytest.mark.parametrize("cat_threshold", [K, K + 1])
+def test_mat_transpose_vec(df, cat_threshold):
+    rng = np.random.default_rng(1234)
+    vec = rng.normal(size=N)
+    mat = tm.from_pandas(df, cat_threshold=cat_threshold)
+    np.testing.assert_equal(mat.transpose_matvec(vec), mat.transpose_matvec(vec))