Merge branch 'main' into add-datasets-ci-tests

adap · Sep 20, 2023 · 59a7b3e · 59a7b3e
2 parents 56c6a3b + 0439e93
commit 59a7b3e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 16 deletions.
diff --git a/datasets/doc/source/ref-api-flwr-datasets.rst b/datasets/doc/source/ref-api-flwr-datasets.rst
@@ -1,39 +1,27 @@
 flwr\_datasets (Python API reference)
 ======================
 
-federated\_dataset
-------------------
-
-.. automodule:: flwr_datasets.federated_dataset
-
 Federated Dataset
 -----------------
 .. autoclass:: flwr_datasets.federated_dataset.FederatedDataset
    :members:
 
+
 partitioner
 -----------
 
 .. automodule:: flwr_datasets.partitioner
 
-partitioner.partitioner
------------------------
-
-.. automodule:: flwr_datasets.partitioner.partitioner
 
 Partitioner
 -----------
 
-.. autoclass:: flwr_datasets.partitioner.partitioner.Partitioner
+.. autoclass:: flwr_datasets.partitioner.Partitioner
    :members:
 
-partitioner.iid_partitioner
----------------------------
-
-.. automodule:: flwr_datasets.partitioner.iid_partitioner
 
 IID Partitioner
 ---------------
 
-.. autoclass:: flwr_datasets.partitioner.iid_partitioner.IidPartitioner
+.. autoclass:: flwr_datasets.partitioner.IidPartitioner
    :members:
diff --git a/datasets/flwr_datasets/partitioner/iid_partitioner.py b/datasets/flwr_datasets/partitioner/iid_partitioner.py
@@ -35,7 +35,18 @@ def __init__(self, num_partitions: int) -> None:
         self._num_partitions = num_partitions
 
     def load_partition(self, idx: int) -> datasets.Dataset:
-        """Load a single IID partition based on the partition index."""
+        """Load a single IID partition based on the partition index.
+
+        Parameters
+        ----------
+        idx: int
+            the index that corresponds to the requested partition
+
+        Returns
+        -------
+        dataset_partition: Dataset
+            single dataset partition
+        """
         return self.dataset.shard(
             num_shards=self._num_partitions, index=idx, contiguous=True
         )