Make DirichletPartitioner available on the package level

adap · Jan 16, 2024 · 36a6a50 · 36a6a50
1 parent 9a2c13b
commit 36a6a50
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py
@@ -15,6 +15,7 @@
 """Flower Datasets Partitioner package."""
 
 
+from .dirichlet_partitioner import DirichletPartitioner
 from .exponential_partitioner import ExponentialPartitioner
 from .iid_partitioner import IidPartitioner
 from .linear_partitioner import LinearPartitioner
@@ -27,6 +28,7 @@
     "IidPartitioner",
     "Partitioner",
     "NaturalIdPartitioner",
+    "DirichletPartitioner",
     "SizePartitioner",
     "LinearPartitioner",
     "SquarePartitioner",