Sklearn type transformer should be automatically loaded with import flytekit

flytekit/__init__.py

@@ -184,7 +184,7 @@
 from flytekit.core.workflow import ImperativeWorkflow as Workflow
 from flytekit.core.workflow import WorkflowFailurePolicy, reference_workflow, workflow
 from flytekit.deck import Deck
-from flytekit.extras import pytorch, tensorflow
+from flytekit.extras import pytorch, sklearn, tensorflow
 from flytekit.extras.persistence import GCSPersistence, HttpPersistence, S3Persistence
 from flytekit.loggers import logger
 from flytekit.models.common import Annotations, AuthRole, Labels

flytekit/extras/pytorch/__init__.py

@@ -1,13 +1,14 @@
 """
-Flytekit PyTorch
-=========================================
 .. currentmodule:: flytekit.extras.pytorch
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
 
     PyTorchCheckpoint
+    PyTorchCheckpointTransformer
+    PyTorchModuleTransformer
+    PyTorchTensorTransformer
 """
 from flytekit.loggers import logger
 

flytekit/extras/sklearn/__init__.py

@@ -1,11 +1,11 @@
 """
-Flytekit Sklearn
-=========================================
 .. currentmodule:: flytekit.extras.sklearn
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
+
+   SklearnEstimatorTransformer
 """
 from flytekit.loggers import logger
 

flytekit/extras/tensorflow/__init__.py

@@ -1,3 +1,14 @@
+"""
+.. currentmodule:: flytekit.extras.tensorflow
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+   TensorFlowRecordFileTransformer
+   TensorFlowRecordsDirTransformer
+"""
+
 from flytekit.loggers import logger
 
 # TODO: abstract this out so that there's an established pattern for registering plugins