Add get_n_splits to KFold classes for compat

Author: kbattocchi

econml/sklearn_extensions/model_selection.py

@@ -137,6 +137,27 @@ def split(self, X, y, sample_weight=None):
         """
         return _split_weighted_sample(self, X, y, sample_weight, is_stratified=False)
 
+    def get_n_splits(self, X, y, groups=None):
+        """Return the number of splitting iterations in the cross-validator.
+
+        Parameters
+        ----------
+        X : object
+            Always ignored, exists for compatibility.
+
+        y : object
+            Always ignored, exists for compatibility.
+
+        groups : object
+            Always ignored, exists for compatibility.
+
+        Returns
+        -------
+        n_splits : int
+            Returns the number of splitting iterations in the cross-validator.
+        """
+        return self.n_splits
+
     def _get_folds_from_splits(self, splits, sample_size):
         folds = []
         sample_indices = np.arange(sample_size)
@@ -213,6 +234,27 @@ def split(self, X, y, sample_weight=None):
         """
         return _split_weighted_sample(self, X, y, sample_weight, is_stratified=True)
 
+    def get_n_splits(self, X, y, groups=None):
+        """Return the number of splitting iterations in the cross-validator.
+
+        Parameters
+        ----------
+        X : object
+            Always ignored, exists for compatibility.
+
+        y : object
+            Always ignored, exists for compatibility.
+
+        groups : object
+            Always ignored, exists for compatibility.
+
+        Returns
+        -------
+        n_splits : int
+            Returns the number of splitting iterations in the cross-validator.
+        """
+        return self.n_splits
+
 
 class GridSearchCVList(BaseEstimator):
     """ An extension of GridSearchCV that allows for passing a list of estimators each with their own