From 2501830983bf884a200ce7b8e25b348d6421b6b7 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Wed, 17 Jul 2024 14:41:45 +0100
Subject: [PATCH] Add numpy import to docstring

---
 datasets/flwr_datasets/partitioner/distribution_partitioner.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/datasets/flwr_datasets/partitioner/distribution_partitioner.py b/datasets/flwr_datasets/partitioner/distribution_partitioner.py
index da540f236510..e01740091b34 100644
--- a/datasets/flwr_datasets/partitioner/distribution_partitioner.py
+++ b/datasets/flwr_datasets/partitioner/distribution_partitioner.py
@@ -92,6 +92,7 @@ class label distribution. `float` values are rounded to the nearest `int`.
     >>> from flwr_datasets import FederatedDataset
     >>> from flwr_datasets.partitioner import DistributionPartitioner
     >>> from pprint import pprint
+    >>> import numpy as np
     >>>
     >>> num_clients = 1_000
     >>> num_unique_labels_per_client = 2