ENH do not create 1d memoryviews of X_binned

- See scikit-learn#17299
lorentzenchr · Jan 5, 2024 · b8d7071 · b8d7071
1 parent 678e399
commit b8d7071
Showing 1 changed file with 26 additions and 26 deletions.
diff --git a/sklearn/ensemble/_hist_gradient_boosting/histogram.pyx b/sklearn/ensemble/_hist_gradient_boosting/histogram.pyx
@@ -187,8 +187,8 @@ cdef class HistogramBuilder:
 
         cdef:
             unsigned int n_samples = sample_indices.shape[0]
-            const X_BINNED_DTYPE_C [::1] X_binned = \
-                self.X_binned[:, feature_idx]
+            const X_BINNED_DTYPE_C [::1, :] X_binned = \
+                self.X_binned
             unsigned int root_node = X_binned.shape[0] == n_samples
             G_H_DTYPE_C [::1] ordered_gradients = \
                 self.ordered_gradients[:n_samples]
@@ -325,7 +325,7 @@ cpdef void _subtract_histograms(
 cpdef void _build_histogram(
         const int feature_idx,
         const unsigned int [::1] sample_indices,  # IN
-        const X_BINNED_DTYPE_C [::1] binned_feature,  # IN
+        const X_BINNED_DTYPE_C [::1, :] binned_feature,  # IN
         const G_H_DTYPE_C [::1] ordered_gradients,  # IN
         const G_H_DTYPE_C [::1] ordered_hessians,  # IN
         hist_struct [:, ::1] out) noexcept nogil:  # OUT
@@ -342,10 +342,10 @@ cpdef void _build_histogram(
         unsigned int bin_idx
 
     for i in range(0, unrolled_upper, 4):
-        bin_0 = binned_feature[sample_indices[i]]
-        bin_1 = binned_feature[sample_indices[i + 1]]
-        bin_2 = binned_feature[sample_indices[i + 2]]
-        bin_3 = binned_feature[sample_indices[i + 3]]
+        bin_0 = binned_feature[sample_indices[i], feature_idx]
+        bin_1 = binned_feature[sample_indices[i + 1], feature_idx]
+        bin_2 = binned_feature[sample_indices[i + 2], feature_idx]
+        bin_3 = binned_feature[sample_indices[i + 3], feature_idx]
 
         out[feature_idx, bin_0].sum_gradients += ordered_gradients[i]
         out[feature_idx, bin_1].sum_gradients += ordered_gradients[i + 1]
@@ -363,7 +363,7 @@ cpdef void _build_histogram(
         out[feature_idx, bin_3].count += 1
 
     for i in range(unrolled_upper, n_node_samples):
-        bin_idx = binned_feature[sample_indices[i]]
+        bin_idx = binned_feature[sample_indices[i], feature_idx]
         out[feature_idx, bin_idx].sum_gradients += ordered_gradients[i]
         out[feature_idx, bin_idx].sum_hessians += ordered_hessians[i]
         out[feature_idx, bin_idx].count += 1
@@ -372,7 +372,7 @@ cpdef void _build_histogram(
 cpdef void _build_histogram_no_hessian(
         const int feature_idx,
         const unsigned int [::1] sample_indices,  # IN
-        const X_BINNED_DTYPE_C [::1] binned_feature,  # IN
+        const X_BINNED_DTYPE_C [::1, :] binned_feature,  # IN
         const G_H_DTYPE_C [::1] ordered_gradients,  # IN
         hist_struct [:, ::1] out) noexcept nogil:  # OUT
     """Return histogram for a given feature, not updating hessians.
@@ -392,10 +392,10 @@ cpdef void _build_histogram_no_hessian(
         unsigned int bin_idx
 
     for i in range(0, unrolled_upper, 4):
-        bin_0 = binned_feature[sample_indices[i]]
-        bin_1 = binned_feature[sample_indices[i + 1]]
-        bin_2 = binned_feature[sample_indices[i + 2]]
-        bin_3 = binned_feature[sample_indices[i + 3]]
+        bin_0 = binned_feature[sample_indices[i], feature_idx]
+        bin_1 = binned_feature[sample_indices[i + 1], feature_idx]
+        bin_2 = binned_feature[sample_indices[i + 2], feature_idx]
+        bin_3 = binned_feature[sample_indices[i + 3], feature_idx]
 
         out[feature_idx, bin_0].sum_gradients += ordered_gradients[i]
         out[feature_idx, bin_1].sum_gradients += ordered_gradients[i + 1]
@@ -408,14 +408,14 @@ cpdef void _build_histogram_no_hessian(
         out[feature_idx, bin_3].count += 1
 
     for i in range(unrolled_upper, n_node_samples):
-        bin_idx = binned_feature[sample_indices[i]]
+        bin_idx = binned_feature[sample_indices[i], feature_idx]
         out[feature_idx, bin_idx].sum_gradients += ordered_gradients[i]
         out[feature_idx, bin_idx].count += 1
 
 
 cpdef void _build_histogram_root(
         const int feature_idx,
-        const X_BINNED_DTYPE_C [::1] binned_feature,  # IN
+        const X_BINNED_DTYPE_C [::1, :] binned_feature,  # IN
         const G_H_DTYPE_C [::1] all_gradients,  # IN
         const G_H_DTYPE_C [::1] all_hessians,  # IN
         hist_struct [:, ::1] out) noexcept nogil:  # OUT
@@ -439,10 +439,10 @@ cpdef void _build_histogram_root(
 
     for i in range(0, unrolled_upper, 4):
 
-        bin_0 = binned_feature[i]
-        bin_1 = binned_feature[i + 1]
-        bin_2 = binned_feature[i + 2]
-        bin_3 = binned_feature[i + 3]
+        bin_0 = binned_feature[i, feature_idx]
+        bin_1 = binned_feature[i + 1, feature_idx]
+        bin_2 = binned_feature[i + 2, feature_idx]
+        bin_3 = binned_feature[i + 3, feature_idx]
 
         out[feature_idx, bin_0].sum_gradients += all_gradients[i]
         out[feature_idx, bin_1].sum_gradients += all_gradients[i + 1]
@@ -460,15 +460,15 @@ cpdef void _build_histogram_root(
         out[feature_idx, bin_3].count += 1
 
     for i in range(unrolled_upper, n_samples):
-        bin_idx = binned_feature[i]
+        bin_idx = binned_feature[i, feature_idx]
         out[feature_idx, bin_idx].sum_gradients += all_gradients[i]
         out[feature_idx, bin_idx].sum_hessians += all_hessians[i]
         out[feature_idx, bin_idx].count += 1
 
 
 cpdef void _build_histogram_root_no_hessian(
         const int feature_idx,
-        const X_BINNED_DTYPE_C [::1] binned_feature,  # IN
+        const X_BINNED_DTYPE_C [::1, :] binned_feature,  # IN
         const G_H_DTYPE_C [::1] all_gradients,  # IN
         hist_struct [:, ::1] out) noexcept nogil:  # OUT
     """Compute histogram of the root node, not updating hessians.
@@ -488,10 +488,10 @@ cpdef void _build_histogram_root_no_hessian(
         unsigned int bin_idx
 
     for i in range(0, unrolled_upper, 4):
-        bin_0 = binned_feature[i]
-        bin_1 = binned_feature[i + 1]
-        bin_2 = binned_feature[i + 2]
-        bin_3 = binned_feature[i + 3]
+        bin_0 = binned_feature[i, feature_idx]
+        bin_1 = binned_feature[i + 1, feature_idx]
+        bin_2 = binned_feature[i + 2, feature_idx]
+        bin_3 = binned_feature[i + 3, feature_idx]
 
         out[feature_idx, bin_0].sum_gradients += all_gradients[i]
         out[feature_idx, bin_1].sum_gradients += all_gradients[i + 1]
@@ -504,6 +504,6 @@ cpdef void _build_histogram_root_no_hessian(
         out[feature_idx, bin_3].count += 1
 
     for i in range(unrolled_upper, n_samples):
-        bin_idx = binned_feature[i]
+        bin_idx = binned_feature[i, feature_idx]
         out[feature_idx, bin_idx].sum_gradients += all_gradients[i]
         out[feature_idx, bin_idx].count += 1