Merge branch 'fds-rename-size-partitioner' into fds-add-size-partitioner

adap · Aug 30, 2024 · e04f840 · e04f840
2 parents fd0cc0a + 5447b9b
commit e04f840
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 8 deletions.
diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py
@@ -19,28 +19,28 @@
 from .distribution_partitioner import DistributionPartitioner
 from .exponential_partitioner import ExponentialPartitioner
 from .grouped_natural_id_partitioner import GroupedNaturalIdPartitioner
+from .id_to_size_function_partitioner import IdToSizeFunctionPartitioner
 from .iid_partitioner import IidPartitioner
 from .inner_dirichlet_partitioner import InnerDirichletPartitioner
 from .linear_partitioner import LinearPartitioner
 from .natural_id_partitioner import NaturalIdPartitioner
 from .partitioner import Partitioner
 from .pathological_partitioner import PathologicalPartitioner
 from .shard_partitioner import ShardPartitioner
-from .size_partitioner import SizePartitioner
 from .square_partitioner import SquarePartitioner
 
 __all__ = [
     "DirichletPartitioner",
     "DistributionPartitioner",
     "ExponentialPartitioner",
     "GroupedNaturalIdPartitioner",
+    "IdToSizeFunctionPartitioner",
     "IidPartitioner",
     "InnerDirichletPartitioner",
     "LinearPartitioner",
     "NaturalIdPartitioner",
     "Partitioner",
     "PathologicalPartitioner",
     "ShardPartitioner",
-    "SizePartitioner",
     "SquarePartitioner",
 ]
diff --git a/datasets/flwr_datasets/partitioner/exponential_partitioner.py b/datasets/flwr_datasets/partitioner/exponential_partitioner.py
@@ -17,10 +17,12 @@
 
 import numpy as np
 
-from flwr_datasets.partitioner.size_partitioner import SizePartitioner
+from flwr_datasets.partitioner.id_to_size_function_partitioner import (
+    IdToSizeFunctionPartitioner,
+)
 
 
-class ExponentialPartitioner(SizePartitioner):
+class ExponentialPartitioner(IdToSizeFunctionPartitioner):
     """Partitioner creates partitions of size that are correlated with exp(id).
 
     The amount of data each client gets is correlated with the exponent of partition ID.

diff --git a/datasets/flwr_datasets/partitioner/linear_partitioner.py b/datasets/flwr_datasets/partitioner/linear_partitioner.py
@@ -15,10 +15,12 @@
 """LinearPartitioner class."""
 
 
-from flwr_datasets.partitioner.size_partitioner import SizePartitioner
+from flwr_datasets.partitioner.id_to_size_function_partitioner import (
+    IdToSizeFunctionPartitioner,
+)
 
 
-class LinearPartitioner(SizePartitioner):
+class LinearPartitioner(IdToSizeFunctionPartitioner):
     """Partitioner creates partitions of size that are linearly correlated with id.
 
     The amount of data each client gets is linearly correlated with the partition ID.

diff --git a/datasets/flwr_datasets/partitioner/square_partitioner.py b/datasets/flwr_datasets/partitioner/square_partitioner.py
@@ -17,10 +17,12 @@
 
 import numpy as np
 
-from flwr_datasets.partitioner.size_partitioner import SizePartitioner
+from flwr_datasets.partitioner.id_to_size_function_partitioner import (
+    IdToSizeFunctionPartitioner,
+)
 
 
-class SquarePartitioner(SizePartitioner):
+class SquarePartitioner(IdToSizeFunctionPartitioner):
     """Partitioner creates partitions of size that are correlated with squared id.
 
     The amount of data each client gets is correlated with the squared partition ID.