ENH Adds n_features_in_ checks to impute module (scikit-learn#18580)

thomasjpfan · web-flow · commit 5fb02bb7bfed · 2020-11-20T14:29:03.000-08:00
* ENH Adds n_features_in_ checks to impute module
diff --git a/sklearn/impute/_base.py b/sklearn/impute/_base.py
@@ -793,16 +793,12 @@ def transform(self, X):
         # Need not validate X again as it would have already been validated
         # in the Imputer calling MissingIndicator
         if not self._precomputed:
-            X = self._validate_input(X, in_fit=True)
+            X = self._validate_input(X, in_fit=False)
         else:
             if not (hasattr(X, 'dtype') and X.dtype.kind == 'b'):
                 raise ValueError("precomputed is True but the input data is "
                                  "not a mask")
 
-        if X.shape[1] != self._n_features:
-            raise ValueError("X has a different number of features "
-                             "than during fitting.")
-
         imputer_mask, features = self._get_missing_features_info(X)
 
         if self.features == "missing-only":
diff --git a/sklearn/impute/_iterative.py b/sklearn/impute/_iterative.py
@@ -474,7 +474,7 @@ def _get_abs_corr_mat(self, X_filled, tolerance=1e-6):
         abs_corr_mat = normalize(abs_corr_mat, norm='l1', axis=0, copy=False)
         return abs_corr_mat
 
-    def _initial_imputation(self, X):
+    def _initial_imputation(self, X, in_fit=False):
         """Perform initial imputation for input X.
 
         Parameters
@@ -483,6 +483,9 @@ def _initial_imputation(self, X):
             Input data, where "n_samples" is the number of samples and
             "n_features" is the number of features.
 
+        in_fit : bool, default=False
+            Whether function is called in fit.
+
         Returns
         -------
         Xt : ndarray, shape (n_samples, n_features)
@@ -506,7 +509,7 @@ def _initial_imputation(self, X):
         else:
             force_all_finite = True
 
-        X = self._validate_data(X, dtype=FLOAT_DTYPES, order="F",
+        X = self._validate_data(X, dtype=FLOAT_DTYPES, order="F", reset=in_fit,
                                 force_all_finite=force_all_finite)
         _check_inputs_dtype(X, self.missing_values)
 
@@ -600,7 +603,8 @@ def fit_transform(self, X, y=None):
 
         self.initial_imputer_ = None
 
-        X, Xt, mask_missing_values, complete_mask = self._initial_imputation(X)
+        X, Xt, mask_missing_values, complete_mask = (
+            self._initial_imputation(X, in_fit=True))
 
         super()._fit_indicator(complete_mask)
         X_indicator = super()._transform_indicator(complete_mask)
diff --git a/sklearn/impute/_knn.py b/sklearn/impute/_knn.py
@@ -10,7 +10,6 @@
 from ..metrics.pairwise import _NAN_METRICS
 from ..neighbors._base import _get_weights
 from ..neighbors._base import _check_weights
-from ..utils import check_array
 from ..utils import is_scalar_nan
 from ..utils._mask import _get_mask
 from ..utils.validation import check_is_fitted
@@ -213,12 +212,9 @@ def transform(self, X):
             force_all_finite = True
         else:
             force_all_finite = "allow-nan"
-        X = check_array(X, accept_sparse=False, dtype=FLOAT_DTYPES,
-                        force_all_finite=force_all_finite, copy=self.copy)
-
-        if X.shape[1] != self._fit_X.shape[1]:
-            raise ValueError("Incompatible dimension between the fitted "
-                             "dataset and the one to be transformed")
+        X = self._validate_data(X, accept_sparse=False, dtype=FLOAT_DTYPES,
+                                force_all_finite=force_all_finite,
+                                copy=self.copy, reset=False)
 
         mask = _get_mask(X, self.missing_values)
         mask_fit_X = self._mask_fit_X
diff --git a/sklearn/tests/test_common.py b/sklearn/tests/test_common.py
@@ -344,7 +344,6 @@ def test_search_cv(estimator, check, request):
     'feature_extraction',
     'feature_selection',
     'gaussian_process',
-    'impute',
     'isotonic',
     'linear_model',
     'manifold',
