feat: Add a predict_proba method on SKLearnClassifier

keras/src/wrappers/sklearn_test.py

@@ -57,7 +57,7 @@ def patched_more_tags(self):
     return parametrize_with_checks(estimators)
 
 
-def dynamic_model(X, y, loss, layers=[10]):
+def dynamic_model(X, y, loss, out_activation_function="softmax", layers=[10]):
     """Creates a basic MLP classifier dynamically choosing binary/multiclass
     classification loss and ouput activations.
     """
@@ -69,7 +69,7 @@ def dynamic_model(X, y, loss, layers=[10]):
         hidden = Dense(layer_size, activation="relu")(hidden)
 
     n_outputs = y.shape[1] if len(y.shape) > 1 else 1
-    out = [Dense(n_outputs, activation="softmax")(hidden)]
+    out = [Dense(n_outputs, activation=out_activation_function)(hidden)]
     model = Model(inp, out)
     model.compile(loss=loss, optimizer="rmsprop")
 
@@ -158,3 +158,54 @@ def test_sklearn_estimator_checks(estimator, check):
             pytest.xfail("Backend not implemented")
         else:
             raise
+
+
+@pytest.mark.parametrize(
+    "estimator, has_predict_proba",
+    [
+        (
+            SKLearnClassifier(
+                model=dynamic_model,
+                model_kwargs={
+                    "out_activation_function": "softmax",
+                    "loss": "categorical_crossentropy",
+                },
+                fit_kwargs={"epochs": 1},
+            ),
+            True,
+        ),
+        (
+            SKLearnClassifier(
+                model=dynamic_model,
+                model_kwargs={
+                    "out_activation_function": "linear",
+                    "loss": "categorical_crossentropy",
+                },
+                fit_kwargs={"epochs": 1},
+            ),
+            False,
+        ),
+    ],
+)
+def test_sklearn_estimator_predict_proba(estimator, has_predict_proba):
+    """Checks that ``SKLearnClassifier`` exposes the ``predict_proba`` method
+    only when the model outputs probabilities.
+    """
+    try:
+        X, y = sklearn.datasets.make_classification(
+            n_samples=100,
+            n_features=10,
+            n_informative=4,
+            n_classes=4,
+            random_state=42,
+        )
+        estimator.fit(X, y)
+        assert hasattr(estimator, "predict_proba") == has_predict_proba
+    except Exception as exc:
+        if keras.config.backend() in ["numpy", "openvino"] and (
+            isinstance(exc, NotImplementedError)
+            or "NotImplementedError" in str(exc)
+        ):
+            pytest.xfail("Backend not implemented")
+        else:
+            raise

keras/src/wrappers/sklearn_wrapper.py

@@ -10,6 +10,7 @@
 from keras.src.wrappers.fixes import type_of_target
 from keras.src.wrappers.utils import TargetReshaper
 from keras.src.wrappers.utils import _check_model
+from keras.src.wrappers.utils import _estimator_has
 from keras.src.wrappers.utils import assert_sklearn_installed
 
 try:
@@ -278,6 +279,15 @@ def dynamic_model(X, y, loss, layers=[10]):
     ```
     """
 
+    @sklearn.utils.metaestimators.available_if(_estimator_has("predict_proba"))
+    def predict_proba(self, X):
+        """Predict class probabilities of the input samples X."""
+        from sklearn.utils.validation import check_is_fitted
+
+        check_is_fitted(self)
+        X = _validate_data(self, X, reset=False)
+        return self.model_.predict(X)
+
     def _process_target(self, y, reset=False):
         """Classifiers do OHE."""
         target_type = type_of_target(y, raise_unknown=True)

keras/src/wrappers/utils.py

@@ -32,6 +32,18 @@ def _check_model(model):
         )
 
 
+def _estimator_has(attr):
+    def check(self):
+        from sklearn.utils.validation import check_is_fitted
+
+        check_is_fitted(self)
+        return (
+            self.model_.layers[-1].activation.__name__ in ("sigmoid", "softmax")
+        )
+
+    return check
+
+
 class TargetReshaper(TransformerMixin, BaseEstimator):
     """Convert 1D targets to 2D and back.