ENH use utility _check_sample_weight in _BaseDiscreteNB (#16263)

Batalex · web-flow · commit 84628b0f4593 · 2020-01-29T11:48:21.000+01:00
diff --git a/sklearn/naive_bayes.py b/sklearn/naive_bayes.py
@@ -569,8 +569,9 @@ def partial_fit(self, X, y, classes=None, sample_weight=None):
         # We convert it to np.float64 to support sample_weight consistently
         Y = Y.astype(np.float64, copy=False)
         if sample_weight is not None:
+            sample_weight = _check_sample_weight(sample_weight, X)
             sample_weight = np.atleast_2d(sample_weight)
-            Y *= check_array(sample_weight).T
+            Y *= sample_weight.T
 
         class_prior = self.class_prior
 
@@ -621,9 +622,9 @@ def fit(self, X, y, sample_weight=None):
         # this means we also don't have to cast X to floating point
         if sample_weight is not None:
             Y = Y.astype(np.float64, copy=False)
-            sample_weight = np.asarray(sample_weight)
+            sample_weight = _check_sample_weight(sample_weight, X)
             sample_weight = np.atleast_2d(sample_weight)
-            Y *= check_array(sample_weight).T
+            Y *= sample_weight.T
 
         class_prior = self.class_prior
 
