Expose functionality via __all__

adap · Nov 15, 2023 · f59fb9c · f59fb9c
1 parent eded53c
commit f59fb9c
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/datasets/flwr_datasets/__init__.py b/datasets/flwr_datasets/__init__.py
@@ -15,10 +15,14 @@
 """Flower Datasets main package."""
 
 
+from flwr_datasets import partitioner, resplitter
 from flwr_datasets.common.version import package_version as _package_version
+from flwr_datasets.federated_dataset import FederatedDataset
 
-from .federated_dataset import FederatedDataset
-
-__all__ = ["FederatedDataset"]
+__all__ = [
+    "FederatedDataset",
+    "partitioner",
+    "resplitter",
+]
 
 __version__ = _package_version