From 0439e936fc917e5758799e048b389340d4241082 Mon Sep 17 00:00:00 2001
From: Adam Narozniak <51029327+adam-narozniak@users.noreply.github.com>
Date: Wed, 20 Sep 2023 11:35:10 +0200
Subject: [PATCH] Fix the reference API documentation (#2397)

---
 datasets/doc/source/ref-api-flwr-datasets.rst  | 18 +++---------------
 .../partitioner/iid_partitioner.py             | 13 ++++++++++++-
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/datasets/doc/source/ref-api-flwr-datasets.rst b/datasets/doc/source/ref-api-flwr-datasets.rst
index 31a324231191..2e6a9e731add 100644
--- 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
index b109e9565e94..c8dbf8294fec 100644
--- 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
         )