Apply suggestions from code review

Co-authored-by: Javier <[email protected]>
adap · Feb 20, 2024 · 5a3f45c · 5a3f45c
1 parent ae4031d
commit 5a3f45c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/datasets/flwr_datasets/partitioner/dirichlet_partitioner.py b/datasets/flwr_datasets/partitioner/dirichlet_partitioner.py
@@ -158,23 +158,23 @@ def _initialize_alpha(
         elif isinstance(alpha, List):
             if len(alpha) != self._num_partitions:
                 raise ValueError(
-                    "The alpha parameter needs to be of length of equal to the "
+                    "If passing alpha as a List, it needs to be of length of equal to "
                     "num_partitions."
                 )
             alpha = np.asarray(alpha)
         elif isinstance(alpha, np.ndarray):
             # pylint: disable=R1720
             if alpha.ndim == 1 and alpha.shape[0] != self._num_partitions:
                 raise ValueError(
-                    "The alpha parameter needs to be of length of equal to"
-                    "the num_partitions."
+                    "If passing alpha as an NDArray, its length needs to be of length "
+                    "equal to num_partitions."
                 )
             elif alpha.ndim == 2:
                 alpha = alpha.flatten()
                 if alpha.shape[0] != self._num_partitions:
                     raise ValueError(
-                        "The alpha parameter needs to be of length of equal to "
-                        "the num_partitions."
+                        "If passing alpha as an NDArray, its size needs to be of length"
+                        " equal to num_partitions."
                     )
         else:
             raise ValueError("The given alpha format is not supported.")

diff --git a/datasets/flwr_datasets/partitioner/dirichlet_partitioner_test.py b/datasets/flwr_datasets/partitioner/dirichlet_partitioner_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 """Test DirichletPartitioner."""
+
+
 # pylint: disable=W0212
 import unittest
 from typing import Tuple, Union