diff --git a/stubs/sklearn/utils/_set_output.pyi b/stubs/sklearn/utils/_set_output.pyi index 234e7fc3..e9f6b919 100644 --- a/stubs/sklearn/utils/_set_output.pyi +++ b/stubs/sklearn/utils/_set_output.pyi @@ -4,4 +4,4 @@ from ..base import BaseEstimator class _SetOutputMixin: def __init_subclass__(cls, auto_wrap_output_keys: None | tuple[str] = ..., **kwargs) -> None: ... - def set_output(self, *, transform: None | Literal["default", "pandas"] = None) -> BaseEstimator: ... + def set_output(self, *, transform: None | Literal["default", "pandas", "polars"] = None) -> BaseEstimator: ...