flwr-nightly 1.8.0.dev20240309__py3-none-any.whl → 1.8.0.dev20240311__py3-none-any.whl

flwr/client/mod/secure_aggregation/secaggplus_mod.py

@@ -72,13 +72,14 @@ class SecAggPlusState:
 
     current_stage: str = Stage.UNMASK
 
-    sid: int = 0
+    nid: int = 0
     sample_num: int = 0
     share_num: int = 0
     threshold: int = 0
     clipping_range: float = 0.0
     target_range: int = 0
     mod_range: int = 0
+    max_weight: float = 0.0
 
     # Secret key (sk) and public key (pk)
     sk1: bytes = b""
@@ -173,12 +174,13 @@ def secaggplus_mod(
 
     # Execute
     if state.current_stage == Stage.SETUP:
+        state.nid = msg.metadata.dst_node_id
         res = _setup(state, configs)
     elif state.current_stage == Stage.SHARE_KEYS:
         res = _share_keys(state, configs)
-    elif state.current_stage == Stage.COLLECT_MASKED_INPUT:
+    elif state.current_stage == Stage.COLLECT_MASKED_VECTORS:
         fit = _get_fit_fn(msg, ctxt, call_next)
-        res = _collect_masked_input(state, configs, fit)
+        res = _collect_masked_vectors(state, configs, fit)
     elif state.current_stage == Stage.UNMASK:
         res = _unmask(state, configs)
     else:
@@ -197,7 +199,7 @@ def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
     # Check the existence of Config.STAGE
     if Key.STAGE not in configs:
         raise KeyError(
-            f"The required key '{Key.STAGE}' is missing from the input `named_values`."
+            f"The required key '{Key.STAGE}' is missing from the ConfigsRecord."
         )
 
     # Check the value type of the Config.STAGE
@@ -213,7 +215,7 @@ def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
         if current_stage != Stage.UNMASK:
             log(WARNING, "Restart from the setup stage")
     # If stage is not "setup",
-    # the stage from `named_values` should be the expected next stage
+    # the stage from configs should be the expected next stage
     else:
         stages = Stage.all()
         expected_next_stage = stages[(stages.index(current_stage) + 1) % len(stages)]
@@ -227,11 +229,10 @@ def check_stage(current_stage: str, configs: ConfigsRecord) -> None:
 # pylint: disable-next=too-many-branches
 def check_configs(stage: str, configs: ConfigsRecord) -> None:
     """Check the validity of the configs."""
-    # Check `named_values` for the setup stage
+    # Check configs for the setup stage
     if stage == Stage.SETUP:
         key_type_pairs = [
             (Key.SAMPLE_NUMBER, int),
-            (Key.SECURE_ID, int),
             (Key.SHARE_NUMBER, int),
             (Key.THRESHOLD, int),
             (Key.CLIPPING_RANGE, float),
@@ -242,7 +243,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
             if key not in configs:
                 raise KeyError(
                     f"Stage {Stage.SETUP}: the required key '{key}' is "
-                    "missing from the input `named_values`."
+                    "missing from the ConfigsRecord."
                 )
             # Bool is a subclass of int in Python,
             # so `isinstance(v, int)` will return True even if v is a boolean.
@@ -265,7 +266,7 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
                     f"Stage {Stage.SHARE_KEYS}: "
                     f"the value for the key '{key}' must be a list of two bytes."
                 )
-    elif stage == Stage.COLLECT_MASKED_INPUT:
+    elif stage == Stage.COLLECT_MASKED_VECTORS:
         key_type_pairs = [
             (Key.CIPHERTEXT_LIST, bytes),
             (Key.SOURCE_LIST, int),
@@ -273,9 +274,9 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
         for key, expected_type in key_type_pairs:
             if key not in configs:
                 raise KeyError(
-                    f"Stage {Stage.COLLECT_MASKED_INPUT}: "
+                    f"Stage {Stage.COLLECT_MASKED_VECTORS}: "
                     f"the required key '{key}' is "
-                    "missing from the input `named_values`."
+                    "missing from the ConfigsRecord."
                 )
             if not isinstance(configs[key], list) or any(
                 elm
@@ -284,21 +285,21 @@ def check_configs(stage: str, configs: ConfigsRecord) -> None:
                 if type(elm) is not expected_type
             ):
                 raise TypeError(
-                    f"Stage {Stage.COLLECT_MASKED_INPUT}: "
+                    f"Stage {Stage.COLLECT_MASKED_VECTORS}: "
                     f"the value for the key '{key}' "
                     f"must be of type List[{expected_type.__name__}]"
                 )
     elif stage == Stage.UNMASK:
         key_type_pairs = [
-            (Key.ACTIVE_SECURE_ID_LIST, int),
-            (Key.DEAD_SECURE_ID_LIST, int),
+            (Key.ACTIVE_NODE_ID_LIST, int),
+            (Key.DEAD_NODE_ID_LIST, int),
         ]
         for key, expected_type in key_type_pairs:
             if key not in configs:
                 raise KeyError(
                     f"Stage {Stage.UNMASK}: "
                     f"the required key '{key}' is "
-                    "missing from the input `named_values`."
+                    "missing from the ConfigsRecord."
                 )
             if not isinstance(configs[key], list) or any(
                 elm
@@ -321,20 +322,20 @@ def _setup(
     # Assigning parameter values to object fields
     sec_agg_param_dict = configs
     state.sample_num = cast(int, sec_agg_param_dict[Key.SAMPLE_NUMBER])
-    state.sid = cast(int, sec_agg_param_dict[Key.SECURE_ID])
-    log(INFO, "Client %d: starting stage 0...", state.sid)
+    log(INFO, "Node %d: starting stage 0...", state.nid)
 
     state.share_num = cast(int, sec_agg_param_dict[Key.SHARE_NUMBER])
     state.threshold = cast(int, sec_agg_param_dict[Key.THRESHOLD])
     state.clipping_range = cast(float, sec_agg_param_dict[Key.CLIPPING_RANGE])
     state.target_range = cast(int, sec_agg_param_dict[Key.TARGET_RANGE])
     state.mod_range = cast(int, sec_agg_param_dict[Key.MOD_RANGE])
+    state.max_weight = cast(float, sec_agg_param_dict[Key.MAX_WEIGHT])
 
-    # Dictionaries containing client secure IDs as keys
+    # Dictionaries containing node IDs as keys
     # and their respective secret shares as values.
     state.rd_seed_share_dict = {}
     state.sk1_share_dict = {}
-    # Dictionary containing client secure IDs as keys
+    # Dictionary containing node IDs as keys
     # and their respective shared secrets (with this client) as values.
     state.ss2_dict = {}
 
@@ -346,7 +347,7 @@ def _setup(
 
     state.sk1, state.pk1 = private_key_to_bytes(sk1), public_key_to_bytes(pk1)
     state.sk2, state.pk2 = private_key_to_bytes(sk2), public_key_to_bytes(pk2)
-    log(INFO, "Client %d: stage 0 completes. uploading public keys...", state.sid)
+    log(INFO, "Node %d: stage 0 completes. uploading public keys...", state.nid)
     return {Key.PUBLIC_KEY_1: state.pk1, Key.PUBLIC_KEY_2: state.pk2}
 
 
@@ -356,7 +357,7 @@ def _share_keys(
 ) -> Dict[str, ConfigsRecordValues]:
     named_bytes_tuples = cast(Dict[str, Tuple[bytes, bytes]], configs)
     key_dict = {int(sid): (pk1, pk2) for sid, (pk1, pk2) in named_bytes_tuples.items()}
-    log(INFO, "Client %d: starting stage 1...", state.sid)
+    log(INFO, "Node %d: starting stage 1...", state.nid)
     state.public_keys_dict = key_dict
 
     # Check if the size is larger than threshold
@@ -373,8 +374,8 @@ def _share_keys(
 
     # Check if public keys of this client are correct in the dictionary
     if (
-        state.public_keys_dict[state.sid][0] != state.pk1
-        or state.public_keys_dict[state.sid][1] != state.pk2
+        state.public_keys_dict[state.nid][0] != state.pk1
+        or state.public_keys_dict[state.nid][1] != state.pk2
     ):
         raise ValueError(
             "Own public keys are displayed in dict incorrectly, should not happen!"
@@ -390,35 +391,35 @@ def _share_keys(
     srcs, dsts, ciphertexts = [], [], []
 
     # Distribute shares
-    for idx, (sid, (_, pk2)) in enumerate(state.public_keys_dict.items()):
-        if sid == state.sid:
-            state.rd_seed_share_dict[state.sid] = b_shares[idx]
-            state.sk1_share_dict[state.sid] = sk1_shares[idx]
+    for idx, (nid, (_, pk2)) in enumerate(state.public_keys_dict.items()):
+        if nid == state.nid:
+            state.rd_seed_share_dict[state.nid] = b_shares[idx]
+            state.sk1_share_dict[state.nid] = sk1_shares[idx]
         else:
             shared_key = generate_shared_key(
                 bytes_to_private_key(state.sk2),
                 bytes_to_public_key(pk2),
             )
-            state.ss2_dict[sid] = shared_key
+            state.ss2_dict[nid] = shared_key
             plaintext = share_keys_plaintext_concat(
-                state.sid, sid, b_shares[idx], sk1_shares[idx]
+                state.nid, nid, b_shares[idx], sk1_shares[idx]
             )
             ciphertext = encrypt(shared_key, plaintext)
-            srcs.append(state.sid)
-            dsts.append(sid)
+            srcs.append(state.nid)
+            dsts.append(nid)
             ciphertexts.append(ciphertext)
 
-    log(INFO, "Client %d: stage 1 completes. uploading key shares...", state.sid)
+    log(INFO, "Node %d: stage 1 completes. uploading key shares...", state.nid)
     return {Key.DESTINATION_LIST: dsts, Key.CIPHERTEXT_LIST: ciphertexts}
 
 
 # pylint: disable-next=too-many-locals
-def _collect_masked_input(
+def _collect_masked_vectors(
     state: SecAggPlusState,
     configs: ConfigsRecord,
     fit: Callable[[], FitRes],
 ) -> Dict[str, ConfigsRecordValues]:
-    log(INFO, "Client %d: starting stage 2...", state.sid)
+    log(INFO, "Node %d: starting stage 2...", state.nid)
     available_clients: List[int] = []
     ciphertexts = cast(List[bytes], configs[Key.CIPHERTEXT_LIST])
     srcs = cast(List[int], configs[Key.SOURCE_LIST])
@@ -435,29 +436,45 @@ def _collect_masked_input(
         available_clients.append(src)
         if src != actual_src:
             raise ValueError(
-                f"Client {state.sid}: received ciphertext "
+                f"Node {state.nid}: received ciphertext "
                 f"from {actual_src} instead of {src}."
             )
-        if dst != state.sid:
+        if dst != state.nid:
             raise ValueError(
-                f"Client {state.sid}: received an encrypted message"
-                f"for Client {dst} from Client {src}."
+                f"Node {state.nid}: received an encrypted message"
+                f"for Node {dst} from Node {src}."
             )
         state.rd_seed_share_dict[src] = rd_seed_share
         state.sk1_share_dict[src] = sk1_share
 
     # Fit client
     fit_res = fit()
-    parameters_factor = fit_res.num_examples
+    if len(fit_res.metrics) > 0:
+        log(
+            WARNING,
+            "The metrics in FitRes will not be preserved or sent to the server.",
+        )
+    ratio = fit_res.num_examples / state.max_weight
+    if ratio > 1:
+        log(
+            WARNING,
+            "Potential overflow warning: the provided weight (%s) exceeds the specified"
+            " max_weight (%s). This may lead to overflow issues.",
+            fit_res.num_examples,
+            state.max_weight,
+        )
+    q_ratio = round(ratio * state.target_range)
+    dq_ratio = q_ratio / state.target_range
+
     parameters = parameters_to_ndarrays(fit_res.parameters)
+    parameters = parameters_multiply(parameters, dq_ratio)
 
     # Quantize parameter update (vector)
     quantized_parameters = quantize(
         parameters, state.clipping_range, state.target_range
     )
 
-    quantized_parameters = parameters_multiply(quantized_parameters, parameters_factor)
-    quantized_parameters = factor_combine(parameters_factor, quantized_parameters)
+    quantized_parameters = factor_combine(q_ratio, quantized_parameters)
 
     dimensions_list: List[Tuple[int, ...]] = [a.shape for a in quantized_parameters]
 
@@ -465,14 +482,14 @@ def _collect_masked_input(
     private_mask = pseudo_rand_gen(state.rd_seed, state.mod_range, dimensions_list)
     quantized_parameters = parameters_addition(quantized_parameters, private_mask)
 
-    for client_id in available_clients:
+    for node_id in available_clients:
         # Add pairwise masks
         shared_key = generate_shared_key(
             bytes_to_private_key(state.sk1),
-            bytes_to_public_key(state.public_keys_dict[client_id][0]),
+            bytes_to_public_key(state.public_keys_dict[node_id][0]),
         )
         pairwise_mask = pseudo_rand_gen(shared_key, state.mod_range, dimensions_list)
-        if state.sid > client_id:
+        if state.nid > node_id:
             quantized_parameters = parameters_addition(
                 quantized_parameters, pairwise_mask
             )
@@ -483,7 +500,7 @@ def _collect_masked_input(
 
     # Take mod of final weight update vector and return to server
     quantized_parameters = parameters_mod(quantized_parameters, state.mod_range)
-    log(INFO, "Client %d: stage 2 completes. uploading masked parameters...", state.sid)
+    log(INFO, "Node %d: stage 2 completed, uploading masked parameters...", state.nid)
     return {
         Key.MASKED_PARAMETERS: [ndarray_to_bytes(arr) for arr in quantized_parameters]
     }
@@ -492,20 +509,19 @@ def _collect_masked_input(
 def _unmask(
     state: SecAggPlusState, configs: ConfigsRecord
 ) -> Dict[str, ConfigsRecordValues]:
-    log(INFO, "Client %d: starting stage 3...", state.sid)
+    log(INFO, "Node %d: starting stage 3...", state.nid)
 
-    active_sids = cast(List[int], configs[Key.ACTIVE_SECURE_ID_LIST])
-    dead_sids = cast(List[int], configs[Key.DEAD_SECURE_ID_LIST])
-    # Send private mask seed share for every avaliable client (including itclient)
+    active_nids = cast(List[int], configs[Key.ACTIVE_NODE_ID_LIST])
+    dead_nids = cast(List[int], configs[Key.DEAD_NODE_ID_LIST])
+    # Send private mask seed share for every avaliable client (including itself)
     # Send first private key share for building pairwise mask for every dropped client
-    if len(active_sids) < state.threshold:
+    if len(active_nids) < state.threshold:
         raise ValueError("Available neighbours number smaller than threshold")
 
-    sids, shares = [], []
-    sids += active_sids
-    shares += [state.rd_seed_share_dict[sid] for sid in active_sids]
-    sids += dead_sids
-    shares += [state.sk1_share_dict[sid] for sid in dead_sids]
+    all_nids, shares = [], []
+    all_nids = active_nids + dead_nids
+    shares += [state.rd_seed_share_dict[nid] for nid in active_nids]
+    shares += [state.sk1_share_dict[nid] for nid in dead_nids]
 
-    log(INFO, "Client %d: stage 3 completes. uploading key shares...", state.sid)
-    return {Key.SECURE_ID_LIST: sids, Key.SHARE_LIST: shares}
+    log(INFO, "Node %d: stage 3 completes. uploading key shares...", state.nid)
+    return {Key.NODE_ID_LIST: all_nids, Key.SHARE_LIST: shares}

flwr/common/grpc.py

@@ -15,7 +15,7 @@
 """Utility functions for gRPC."""
 
 
-from logging import INFO
+from logging import DEBUG
 from typing import Optional
 
 import grpc
@@ -49,12 +49,12 @@ def create_channel(
 
     if insecure:
         channel = grpc.insecure_channel(server_address, options=channel_options)
-        log(INFO, "Opened insecure gRPC connection (no certificates were passed)")
+        log(DEBUG, "Opened insecure gRPC connection (no certificates were passed)")
     else:
         ssl_channel_credentials = grpc.ssl_channel_credentials(root_certificates)
         channel = grpc.secure_channel(
             server_address, ssl_channel_credentials, options=channel_options
         )
-        log(INFO, "Opened secure gRPC connection using certificates")
+        log(DEBUG, "Opened secure gRPC connection using certificates")
 
     return channel

flwr/common/logger.py

@@ -18,21 +18,86 @@
 import logging
 from logging import WARN, LogRecord
 from logging.handlers import HTTPHandler
-from typing import Any, Dict, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Optional, TextIO, Tuple
 
 # Create logger
 LOGGER_NAME = "flwr"
 FLOWER_LOGGER = logging.getLogger(LOGGER_NAME)
 FLOWER_LOGGER.setLevel(logging.DEBUG)
 
-DEFAULT_FORMATTER = logging.Formatter(
-    "%(levelname)s %(name)s %(asctime)s | %(filename)s:%(lineno)d | %(message)s"
-)
+LOG_COLORS = {
+    "DEBUG": "\033[94m",  # Blue
+    "INFO": "\033[92m",  # Green
+    "WARNING": "\033[93m",  # Yellow
+    "ERROR": "\033[91m",  # Red
+    "CRITICAL": "\033[95m",  # Magenta
+    "RESET": "\033[0m",  # Reset to default
+}
+
+if TYPE_CHECKING:
+    StreamHandler = logging.StreamHandler[Any]
+else:
+    StreamHandler = logging.StreamHandler
+
+
+class ConsoleHandler(StreamHandler):
+    """Console handler that allows configurable formatting."""
+
+    def __init__(
+        self,
+        timestamps: bool = False,
+        json: bool = False,
+        colored: bool = True,
+        stream: Optional[TextIO] = None,
+    ) -> None:
+        super().__init__(stream)
+        self.timestamps = timestamps
+        self.json = json
+        self.colored = colored
+
+    def emit(self, record: LogRecord) -> None:
+        """Emit a record."""
+        if self.json:
+            record.message = record.getMessage().replace("\t", "").strip()
+
+            # Check if the message is empty
+            if not record.message:
+                return
+
+        super().emit(record)
+
+    def format(self, record: LogRecord) -> str:
+        """Format function that adds colors to log level."""
+        seperator = " " * (8 - len(record.levelname))
+        if self.json:
+            log_fmt = "{lvl='%(levelname)s', time='%(asctime)s', msg='%(message)s'}"
+        else:
+            log_fmt = (
+                f"{LOG_COLORS[record.levelname] if self.colored else ''}"
+                f"%(levelname)s {'%(asctime)s' if self.timestamps else ''}"
+                f"{LOG_COLORS['RESET'] if self.colored else ''}"
+                f": {seperator} %(message)s"
+            )
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+def update_console_handler(level: int, timestamps: bool, colored: bool) -> None:
+    """Update the logging handler."""
+    for handler in logging.getLogger(LOGGER_NAME).handlers:
+        if isinstance(handler, ConsoleHandler):
+            handler.setLevel(level)
+            handler.timestamps = timestamps
+            handler.colored = colored
+
 
 # Configure console logger
-console_handler = logging.StreamHandler()
-console_handler.setLevel(logging.DEBUG)
-console_handler.setFormatter(DEFAULT_FORMATTER)
+console_handler = ConsoleHandler(
+    timestamps=False,
+    json=False,
+    colored=True,
+)
+console_handler.setLevel(logging.INFO)
 FLOWER_LOGGER.addHandler(console_handler)
 
 
@@ -103,11 +168,10 @@ def warn_experimental_feature(name: str) -> None:
     """Warn the user when they use an experimental feature."""
     log(
         WARN,
-        """
-        EXPERIMENTAL FEATURE: %s
+        """EXPERIMENTAL FEATURE: %s
 
-        This is an experimental feature. It could change significantly or be removed
-        entirely in future versions of Flower.
+            This is an experimental feature. It could change significantly or be removed
+            entirely in future versions of Flower.
         """,
         name,
     )
@@ -117,11 +181,10 @@ def warn_deprecated_feature(name: str) -> None:
     """Warn the user when they use a deprecated feature."""
     log(
         WARN,
-        """
-        DEPRECATED FEATURE: %s
+        """DEPRECATED FEATURE: %s
 
-        This is a deprecated feature. It will be removed
-        entirely in future versions of Flower.
+            This is a deprecated feature. It will be removed
+            entirely in future versions of Flower.
         """,
         name,
     )

flwr/common/object_ref.py

@@ -0,0 +1,140 @@
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Helper functions to load objects from a reference."""
+
+
+import ast
+import importlib
+from importlib.util import find_spec
+from typing import Any, Optional, Tuple, Type
+
+OBJECT_REF_HELP_STR = """
+\n\nThe object reference string should have the form <module>:<attribute>. Valid
+examples include `client:app` and `project.package.module:wrapper.app`. It must
+refer to a module on the PYTHONPATH and the module needs to have the specified
+attribute.
+"""
+
+
+def validate(
+    module_attribute_str: str,
+) -> Tuple[bool, Optional[str]]:
+    """Validate object reference.
+
+    The object reference string should have the form <module>:<attribute>. Valid
+    examples include `client:app` and `project.package.module:wrapper.app`. It must
+    refer to a module on the PYTHONPATH and the module needs to have the specified
+    attribute.
+
+    Returns
+    -------
+    Tuple[bool, Optional[str]]
+        A boolean indicating whether an object reference is valid and
+        the reason why it might not be.
+    """
+    module_str, _, attributes_str = module_attribute_str.partition(":")
+    if not module_str:
+        return (
+            False,
+            f"Missing module in {module_attribute_str}{OBJECT_REF_HELP_STR}",
+        )
+    if not attributes_str:
+        return (
+            False,
+            f"Missing attribute in {module_attribute_str}{OBJECT_REF_HELP_STR}",
+        )
+
+    # Load module
+    module = find_spec(module_str)
+    if module and module.origin:
+        if not _find_attribute_in_module(module.origin, attributes_str):
+            return (
+                False,
+                f"Unable to find attribute {attributes_str} in module {module_str}"
+                f"{OBJECT_REF_HELP_STR}",
+            )
+        return (True, None)
+
+    return (
+        False,
+        f"Unable to load module {module_str}{OBJECT_REF_HELP_STR}",
+    )
+
+
+def load_app(
+    module_attribute_str: str,
+    error_type: Type[Exception],
+) -> Any:
+    """Return the object specified in a module attribute string.
+
+    The module/attribute string should have the form <module>:<attribute>. Valid
+    examples include `client:app` and `project.package.module:wrapper.app`. It must
+    refer to a module on the PYTHONPATH, the module needs to have the specified
+    attribute.
+    """
+    valid, error_msg = validate(module_attribute_str)
+    if not valid and error_msg:
+        raise error_type(error_msg) from None
+
+    module_str, _, attributes_str = module_attribute_str.partition(":")
+
+    try:
+        module = importlib.import_module(module_str)
+    except ModuleNotFoundError:
+        raise error_type(
+            f"Unable to load module {module_str}{OBJECT_REF_HELP_STR}",
+        ) from None
+
+    # Recursively load attribute
+    attribute = module
+    try:
+        for attribute_str in attributes_str.split("."):
+            attribute = getattr(attribute, attribute_str)
+    except AttributeError:
+        raise error_type(
+            f"Unable to load attribute {attributes_str} from module {module_str}"
+            f"{OBJECT_REF_HELP_STR}",
+        ) from None
+
+    return attribute
+
+
+def _find_attribute_in_module(file_path: str, attribute_name: str) -> bool:
+    """Check if attribute_name exists in module's abstract symbolic tree."""
+    with open(file_path, encoding="utf-8") as file:
+        node = ast.parse(file.read(), filename=file_path)
+
+    for n in ast.walk(node):
+        if isinstance(n, ast.Assign):
+            for target in n.targets:
+                if isinstance(target, ast.Name) and target.id == attribute_name:
+                    return True
+                if _is_module_in_all(attribute_name, target, n):
+                    return True
+    return False
+
+
+def _is_module_in_all(attribute_name: str, target: ast.expr, n: ast.Assign) -> bool:
+    """Now check if attribute_name is in __all__."""
+    if isinstance(target, ast.Name) and target.id == "__all__":
+        if isinstance(n.value, ast.List):
+            for elt in n.value.elts:
+                if isinstance(elt, ast.Str) and elt.s == attribute_name:
+                    return True
+        elif isinstance(n.value, ast.Tuple):
+            for elt in n.value.elts:
+                if isinstance(elt, ast.Str) and elt.s == attribute_name:
+                    return True
+    return False

flwr/common/secure_aggregation/ndarrays_arithmetic.py

@@ -15,7 +15,7 @@
 """Utility functions for performing operations on Numpy NDArrays."""
 
 
-from typing import Any, List, Tuple
+from typing import Any, List, Tuple, Union
 
 import numpy as np
 from numpy.typing import DTypeLike, NDArray
@@ -68,14 +68,14 @@ def parameters_mod(parameters: List[NDArray[Any]], divisor: int) -> List[NDArray
 
 
 def parameters_multiply(
-    parameters: List[NDArray[Any]], multiplier: int
+    parameters: List[NDArray[Any]], multiplier: Union[int, float]
 ) -> List[NDArray[Any]]:
-    """Multiply parameters by an integer multiplier."""
+    """Multiply parameters by an integer/float multiplier."""
     return [parameters[idx] * multiplier for idx in range(len(parameters))]
 
 
 def parameters_divide(
-    parameters: List[NDArray[Any]], divisor: int
+    parameters: List[NDArray[Any]], divisor: Union[int, float]
 ) -> List[NDArray[Any]]:
-    """Divide weight by an integer divisor."""
+    """Divide weight by an integer/float divisor."""
     return [parameters[idx] / divisor for idx in range(len(parameters))]