Fix docstrings

zillow · Aug 15, 2024 · 04fd72b · 04fd72b
1 parent d324bc0
commit 04fd72b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/quantile_forest/_quantile_forest.py b/quantile_forest/_quantile_forest.py
@@ -246,7 +246,7 @@ def _get_y_train_leaves(self, X, y, sorter=None, sample_weight=None):
         Returns
         -------
         y_train_leaves : array-like of shape \
-                (n_estimators, n_leaves, n_indices, n_outputs)
+                (n_estimators, n_leaves, n_outputs, n_indices)
             List of trees, each with a list of nodes, each with a list of
             indices of the training samples residing at that node. Nodes with
             no samples (e.g., internal nodes) are empty. Internal nodes are
@@ -369,7 +369,7 @@ def _get_y_bound_leaves(self, y, y_train_leaves):
             The target outputs for each y value.
 
         y_train_leaves : array-like of shape \
-                (n_estimators, n_leaves, n_indices, n_outputs)
+                (n_estimators, n_leaves, n_outputs, n_indices)
             List of trees, each with a list of nodes, each with a list of
             indices of the training samples residing at that node. Nodes with
             no samples (e.g., internal nodes) are empty. Internal nodes are

diff --git a/quantile_forest/_quantile_forest_fast.pyx b/quantile_forest/_quantile_forest_fast.pyx
@@ -563,7 +563,7 @@ cdef class QuantileForest:
         Training target values. Assumes values are sorted in ascending order.
 
     y_train_leaves : array-like of shape \
-            (n_estimators, n_leaves, n_indices, n_outputs)
+            (n_estimators, n_leaves, n_outputs, n_indices)
         List of trees, each with a list of nodes, each with a list of indices
         of the training samples residing at that node. Nodes with no samples
         (e.g., internal nodes) are empty. Internal nodes are included so that