flwr-nightly 1.22.0.dev20250913__tar.gz → 1.22.0.dev20250916__tar.gz

{flwr_nightly-1.22.0.dev20250913 → flwr_nightly-1.22.0.dev20250916}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: flwr-nightly
-Version: 1.22.0.dev20250913
+Version: 1.22.0.dev20250916
 Summary: Flower: A Friendly Federated AI Framework
 License: Apache-2.0
 Keywords: Artificial Intelligence,Federated AI,Federated Analytics,Federated Evaluation,Federated Learning,Flower,Machine Learning
@@ -102,25 +102,15 @@ Meet the Flower community on [flower.ai](https://flower.ai)!
 
 Flower's goal is to make federated learning accessible to everyone. This series of tutorials introduces the fundamentals of federated learning and how to implement them in Flower.
 
-0. **What is Federated Learning?**
+0. **[What is Federated Learning?](https://flower.ai/docs/framework/main/en/tutorial-series-what-is-federated-learning.html)**
 
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-what-is-federated-learning.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-what-is-federated-learning.ipynb))
+1. **[An Introduction to Federated Learning](https://flower.ai/docs/framework/main/en/tutorial-series-get-started-with-flower-pytorch.html)**
 
-1. **An Introduction to Federated Learning**
+2. **[Using Strategies in Federated Learning](https://flower.ai/docs/framework/main/en/tutorial-series-use-a-federated-learning-strategy-pytorch.html)**
 
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-get-started-with-flower-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-get-started-with-flower-pytorch.ipynb))
+3. **[Customize a Flower Strategy](https://flower.ai/docs/framework/main/en/tutorial-series-build-a-strategy-from-scratch-pytorch.html)**
 
-2. **Using Strategies in Federated Learning**
-
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-use-a-federated-learning-strategy-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-use-a-federated-learning-strategy-pytorch.ipynb))
-
-3. **Building Strategies for Federated Learning**
-
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-build-a-strategy-from-scratch-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-build-a-strategy-from-scratch-pytorch.ipynb))
-
-4. **Custom Clients for Federated Learning**
-
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-customize-the-client-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-customize-the-client-pytorch.ipynb))
+4. **[Communicate Custom Messages](https://flower.ai/docs/framework/main/en/tutorial-series-customize-the-client-pytorch.html)**
 
 Stay tuned, more tutorials are coming soon. Topics include **Privacy and Security in Federated Learning**, and **Scaling Federated Learning**.
 

{flwr_nightly-1.22.0.dev20250913 → flwr_nightly-1.22.0.dev20250916}/README.md

@@ -46,25 +46,15 @@ Meet the Flower community on [flower.ai](https://flower.ai)!
 
 Flower's goal is to make federated learning accessible to everyone. This series of tutorials introduces the fundamentals of federated learning and how to implement them in Flower.
 
-0. **What is Federated Learning?**
+0. **[What is Federated Learning?](https://flower.ai/docs/framework/main/en/tutorial-series-what-is-federated-learning.html)**
 
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-what-is-federated-learning.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-what-is-federated-learning.ipynb))
+1. **[An Introduction to Federated Learning](https://flower.ai/docs/framework/main/en/tutorial-series-get-started-with-flower-pytorch.html)**
 
-1. **An Introduction to Federated Learning**
+2. **[Using Strategies in Federated Learning](https://flower.ai/docs/framework/main/en/tutorial-series-use-a-federated-learning-strategy-pytorch.html)**
 
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-get-started-with-flower-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-get-started-with-flower-pytorch.ipynb))
+3. **[Customize a Flower Strategy](https://flower.ai/docs/framework/main/en/tutorial-series-build-a-strategy-from-scratch-pytorch.html)**
 
-2. **Using Strategies in Federated Learning**
-
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-use-a-federated-learning-strategy-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-use-a-federated-learning-strategy-pytorch.ipynb))
-
-3. **Building Strategies for Federated Learning**
-
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-build-a-strategy-from-scratch-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-build-a-strategy-from-scratch-pytorch.ipynb))
-
-4. **Custom Clients for Federated Learning**
-
-   [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adap/flower/blob/main/framework/docs/source/tutorial-series-customize-the-client-pytorch.ipynb) (or open the [Jupyter Notebook](https://github.com/adap/flower/blob/main/framework/docs/source/tutorial-series-customize-the-client-pytorch.ipynb))
+4. **[Communicate Custom Messages](https://flower.ai/docs/framework/main/en/tutorial-series-customize-the-client-pytorch.html)**
 
 Stay tuned, more tutorials are coming soon. Topics include **Privacy and Security in Federated Learning**, and **Scaling Federated Learning**.
 

{flwr_nightly-1.22.0.dev20250913 → flwr_nightly-1.22.0.dev20250916}/py/flwr/cli/new/new.py

@@ -35,7 +35,6 @@ class MlFramework(str, Enum):
     """Available frameworks."""
 
     PYTORCH = "PyTorch"
-    PYTORCH_MSG_API = "PyTorch (Message API)"
     TENSORFLOW = "TensorFlow"
     SKLEARN = "sklearn"
     HUGGINGFACE = "HuggingFace"
@@ -44,6 +43,7 @@ class MlFramework(str, Enum):
     NUMPY = "NumPy"
     FLOWERTUNE = "FlowerTune"
     BASELINE = "Flower Baseline"
+    PYTORCH_LEGACY_API = "PyTorch (Legacy API, deprecated)"
 
 
 class LlmChallengeName(str, Enum):
@@ -155,8 +155,8 @@ def new(
     if framework_str == MlFramework.BASELINE:
         framework_str = "baseline"
 
-    if framework_str == MlFramework.PYTORCH_MSG_API:
-        framework_str = "pytorch_msg_api"
+    if framework_str == MlFramework.PYTORCH_LEGACY_API:
+        framework_str = "pytorch_legacy_api"
 
     print(
         typer.style(
@@ -247,14 +247,14 @@ def new(
             MlFramework.TENSORFLOW.value,
             MlFramework.SKLEARN.value,
             MlFramework.NUMPY.value,
-            "pytorch_msg_api",
+            "pytorch_legacy_api",
         ]
         if framework_str in frameworks_with_tasks:
             files[f"{import_name}/task.py"] = {
                 "template": f"app/code/task.{template_name}.py.tpl"
             }
 
-        if framework_str == "pytorch_msg_api":
+        if framework_str == "pytorch_legacy_api":
             # Use custom __init__ that better captures name of framework
             files[f"{import_name}/__init__.py"] = {
                 "template": f"app/code/__init__.{framework_str}.py.tpl"

{flwr_nightly-1.22.0.dev20250913 → flwr_nightly-1.22.0.dev20250916}/py/flwr/serverapp/strategy/__init__.py

@@ -22,6 +22,10 @@ from .dp_fixed_clipping import (
 from .fedadagrad import FedAdagrad
 from .fedadam import FedAdam
 from .fedavg import FedAvg
+from .fedavgm import FedAvgM
+from .fedmedian import FedMedian
+from .fedtrimmedavg import FedTrimmedAvg
+from .fedxgb_bagging import FedXgbBagging
 from .fedyogi import FedYogi
 from .result import Result
 from .strategy import Strategy
@@ -32,6 +36,10 @@ __all__ = [
     "FedAdagrad",
     "FedAdam",
     "FedAvg",
+    "FedAvgM",
+    "FedMedian",
+    "FedTrimmedAvg",
+    "FedXgbBagging",
     "FedYogi",
     "Result",
     "Strategy",

{flwr_nightly-1.22.0.dev20250913 → flwr_nightly-1.22.0.dev20250916}/py/flwr/serverapp/strategy/fedavg.py

@@ -16,7 +16,7 @@
 
 
 from collections.abc import Iterable
-from logging import INFO
+from logging import INFO, WARNING
 from typing import Callable, Optional
 
 from flwr.common import (
@@ -67,7 +67,7 @@ class FedAvg(Strategy):
     arrayrecord_key : str (default: "arrays")
         Key used to store the ArrayRecord when constructing Messages.
     configrecord_key : str (default: "config")
-         Key used to store the ConfigRecord when constructing Messages.
+        Key used to store the ConfigRecord when constructing Messages.
     train_metrics_aggr_fn : Optional[callable] (default: None)
         Function with signature (list[RecordDict], str) -> MetricRecord,
         used to aggregate MetricRecords from training round replies.
@@ -111,6 +111,20 @@ class FedAvg(Strategy):
             evaluate_metrics_aggr_fn or aggregate_metricrecords
         )
 
+        if self.fraction_evaluate == 0.0:
+            self.min_evaluate_nodes = 0
+            log(
+                WARNING,
+                "fraction_evaluate is set to 0.0. "
+                "Federated evaluation will be skipped.",
+            )
+        if self.fraction_train == 0.0:
+            self.min_train_nodes = 0
+            log(
+                WARNING,
+                "fraction_train is set to 0.0. Federated training will be skipped.",
+            )
+
     def summary(self) -> None:
         """Log summary configuration of the strategy."""
         log(INFO, "\t├──> Sampling:")
@@ -150,6 +164,9 @@ class FedAvg(Strategy):
         self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
     ) -> Iterable[Message]:
         """Configure the next round of federated training."""
+        # Do not configure federated train if fraction_train is 0.
+        if self.fraction_train == 0.0:
+            return []
         # Sample nodes
         num_nodes = int(len(list(grid.get_node_ids())) * self.fraction_train)
         sample_size = max(num_nodes, self.min_train_nodes)
@@ -259,6 +276,10 @@ class FedAvg(Strategy):
         self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
     ) -> Iterable[Message]:
         """Configure the next round of federated evaluation."""
+        # Do not configure federated evaluation if fraction_evaluate is 0.
+        if self.fraction_evaluate == 0.0:
+            return []
+
         # Sample nodes
         num_nodes = int(len(list(grid.get_node_ids())) * self.fraction_evaluate)
         sample_size = max(num_nodes, self.min_evaluate_nodes)

flwr_nightly-1.22.0.dev20250916/py/flwr/serverapp/strategy/fedavgm.py

@@ -0,0 +1,198 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Federated Averaging with Momentum (FedAvgM) [Hsu et al., 2019] strategy.
+
+Paper: arxiv.org/pdf/1909.06335.pdf
+"""
+
+
+from collections import OrderedDict
+from collections.abc import Iterable
+from logging import INFO
+from typing import Callable, Optional
+
+from flwr.common import (
+    Array,
+    ArrayRecord,
+    ConfigRecord,
+    Message,
+    MetricRecord,
+    NDArrays,
+    RecordDict,
+    log,
+)
+from flwr.server import Grid
+
+from ..exception import AggregationError
+from .fedavg import FedAvg
+
+
+class FedAvgM(FedAvg):
+    """Federated Averaging with Momentum strategy.
+
+    Implementation based on https://arxiv.org/abs/1909.06335
+
+    Parameters
+    ----------
+    fraction_train : float (default: 1.0)
+        Fraction of nodes used during training. In case `min_train_nodes`
+        is larger than `fraction_train * total_connected_nodes`, `min_train_nodes`
+        will still be sampled.
+    fraction_evaluate : float (default: 1.0)
+        Fraction of nodes used during validation. In case `min_evaluate_nodes`
+        is larger than `fraction_evaluate * total_connected_nodes`,
+        `min_evaluate_nodes` will still be sampled.
+    min_train_nodes : int (default: 2)
+        Minimum number of nodes used during training.
+    min_evaluate_nodes : int (default: 2)
+        Minimum number of nodes used during validation.
+    min_available_nodes : int (default: 2)
+        Minimum number of total nodes in the system.
+    weighted_by_key : str (default: "num-examples")
+        The key within each MetricRecord whose value is used as the weight when
+        computing weighted averages for both ArrayRecords and MetricRecords.
+    arrayrecord_key : str (default: "arrays")
+        Key used to store the ArrayRecord when constructing Messages.
+    configrecord_key : str (default: "config")
+        Key used to store the ConfigRecord when constructing Messages.
+    train_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    evaluate_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    server_learning_rate: float (default: 1.0)
+        Server-side learning rate used in server-side optimization.
+    server_momentum: float (default: 0.0)
+        Server-side momentum factor used for FedAvgM.
+    """
+
+    def __init__(  # pylint: disable=R0913, R0917
+        self,
+        fraction_train: float = 1.0,
+        fraction_evaluate: float = 1.0,
+        min_train_nodes: int = 2,
+        min_evaluate_nodes: int = 2,
+        min_available_nodes: int = 2,
+        weighted_by_key: str = "num-examples",
+        arrayrecord_key: str = "arrays",
+        configrecord_key: str = "config",
+        train_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        evaluate_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        server_learning_rate: float = 1.0,
+        server_momentum: float = 0.0,
+    ) -> None:
+        super().__init__(
+            fraction_train=fraction_train,
+            fraction_evaluate=fraction_evaluate,
+            min_train_nodes=min_train_nodes,
+            min_evaluate_nodes=min_evaluate_nodes,
+            min_available_nodes=min_available_nodes,
+            weighted_by_key=weighted_by_key,
+            arrayrecord_key=arrayrecord_key,
+            configrecord_key=configrecord_key,
+            train_metrics_aggr_fn=train_metrics_aggr_fn,
+            evaluate_metrics_aggr_fn=evaluate_metrics_aggr_fn,
+        )
+        self.server_learning_rate = server_learning_rate
+        self.server_momentum = server_momentum
+        self.server_opt: bool = (self.server_momentum != 0.0) or (
+            self.server_learning_rate != 1.0
+        )
+        self.current_arrays: Optional[ArrayRecord] = None
+        self.momentum_vector: Optional[NDArrays] = None
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        opt_status = "ON" if self.server_opt else "OFF"
+        log(INFO, "\t├──> FedAvgM settings:")
+        log(INFO, "\t|\t├── Server optimization: %s", opt_status)
+        log(INFO, "\t|\t├── Server learning rate: %s", self.server_learning_rate)
+        log(INFO, "\t|\t└── Server Momentum: %s", self.server_momentum)
+        super().summary()
+
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of federated training."""
+        if self.current_arrays is None:
+            self.current_arrays = arrays
+        return super().configure_train(server_round, arrays, config, grid)
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        # Call FedAvg aggregate_train to perform validation and aggregation
+        aggregated_arrays, aggregated_metrics = super().aggregate_train(
+            server_round, replies
+        )
+
+        # following convention described in
+        # https://pytorch.org/docs/stable/generated/torch.optim.SGD.html
+        if self.server_opt and aggregated_arrays is not None:
+            # The initial parameters should be set in `start()` method already
+            if self.current_arrays is None:
+                raise AggregationError(
+                    "No initial parameters set for FedAvgM. "
+                    "Ensure that `configure_train` has been called before aggregation."
+                )
+            ndarrays = self.current_arrays.to_numpy_ndarrays()
+            aggregated_ndarrays = aggregated_arrays.to_numpy_ndarrays()
+
+            # Preserve keys for arrays in ArrayRecord
+            array_keys = list(aggregated_arrays.keys())
+            aggregated_arrays.clear()
+
+            # Remember that updates are the opposite of gradients
+            pseudo_gradient = [
+                old - new for new, old in zip(aggregated_ndarrays, ndarrays)
+            ]
+            if self.server_momentum > 0.0:
+                if self.momentum_vector is None:
+                    # Initialize momentum vector in the first round
+                    self.momentum_vector = pseudo_gradient
+                else:
+                    self.momentum_vector = [
+                        self.server_momentum * mv + pg
+                        for mv, pg in zip(self.momentum_vector, pseudo_gradient)
+                    ]
+
+                # No nesterov for now
+                pseudo_gradient = self.momentum_vector
+
+            # SGD and convert back to ArrayRecord
+            updated_array_list = [
+                Array(old - self.server_learning_rate * pg)
+                for old, pg in zip(ndarrays, pseudo_gradient)
+            ]
+            aggregated_arrays = ArrayRecord(
+                OrderedDict(zip(array_keys, updated_array_list))
+            )
+
+            # Update current weights
+            self.current_arrays = aggregated_arrays
+
+        return aggregated_arrays, aggregated_metrics

flwr_nightly-1.22.0.dev20250916/py/flwr/serverapp/strategy/fedmedian.py

@@ -0,0 +1,71 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Federated Median (FedMedian) [Yin et al., 2018] strategy.
+
+Paper: arxiv.org/pdf/1803.01498v1.pdf
+"""
+
+
+from collections.abc import Iterable
+from typing import Optional, cast
+
+import numpy as np
+
+from flwr.common import Array, ArrayRecord, Message, MetricRecord
+
+from .fedavg import FedAvg
+
+
+class FedMedian(FedAvg):
+    """Federated Median (FedMedian) strategy.
+
+    Implementation based on https://arxiv.org/pdf/1803.01498v1
+    """
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        # Call FedAvg aggregate_train to perform validation and aggregation
+        valid_replies, _ = self._check_and_log_replies(replies, is_train=True)
+
+        if not valid_replies:
+            return None, None
+
+        # Aggregate ArrayRecords using median
+        # Get the key for the only ArrayRecord from the first Message
+        record_key = list(valid_replies[0].content.array_records.keys())[0]
+        # Preserve keys for arrays in ArrayRecord
+        array_keys = list(valid_replies[0].content[record_key].keys())
+
+        # Compute median for each layer and construct ArrayRecord
+        arrays = ArrayRecord()
+        for array_key in array_keys:
+            # Get the corresponding layer from each client
+            layers = [
+                cast(ArrayRecord, msg.content[record_key]).pop(array_key).numpy()
+                for msg in valid_replies
+            ]
+            # Compute median and save as Array in ArrayRecord
+            arrays[array_key] = Array(np.median(np.stack(layers), axis=0))
+
+        # Aggregate MetricRecords
+        metrics = self.train_metrics_aggr_fn(
+            [msg.content for msg in valid_replies],
+            self.weighted_by_key,
+        )
+        return arrays, metrics

flwr_nightly-1.22.0.dev20250916/py/flwr/serverapp/strategy/fedtrimmedavg.py

@@ -0,0 +1,176 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Federated Averaging with Trimmed Mean [Dong Yin, et al., 2021].
+
+Paper: arxiv.org/abs/1803.01498
+"""
+
+
+from collections.abc import Iterable
+from logging import INFO
+from typing import Callable, Optional, cast
+
+import numpy as np
+
+from flwr.common import Array, ArrayRecord, Message, MetricRecord, NDArray, RecordDict
+from flwr.common.logger import log
+
+from ..exception import AggregationError
+from .fedavg import FedAvg
+
+
+class FedTrimmedAvg(FedAvg):
+    """Federated Averaging with Trimmed Mean [Dong Yin, et al., 2021].
+
+    Implemented based on: https://arxiv.org/abs/1803.01498
+
+    Parameters
+    ----------
+    fraction_train : float (default: 1.0)
+        Fraction of nodes used during training. In case `min_train_nodes`
+        is larger than `fraction_train * total_connected_nodes`, `min_train_nodes`
+        will still be sampled.
+    fraction_evaluate : float (default: 1.0)
+        Fraction of nodes used during validation. In case `min_evaluate_nodes`
+        is larger than `fraction_evaluate * total_connected_nodes`,
+        `min_evaluate_nodes` will still be sampled.
+    min_train_nodes : int (default: 2)
+        Minimum number of nodes used during training.
+    min_evaluate_nodes : int (default: 2)
+        Minimum number of nodes used during validation.
+    min_available_nodes : int (default: 2)
+        Minimum number of total nodes in the system.
+    weighted_by_key : str (default: "num-examples")
+        The key within each MetricRecord whose value is used as the weight when
+        computing weighted averages for both ArrayRecords and MetricRecords.
+    arrayrecord_key : str (default: "arrays")
+        Key used to store the ArrayRecord when constructing Messages.
+    configrecord_key : str (default: "config")
+        Key used to store the ConfigRecord when constructing Messages.
+    train_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    evaluate_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    beta : float (default: 0.2)
+        Fraction to cut off of both tails of the distribution.
+    """
+
+    def __init__(  # pylint: disable=R0913, R0917
+        self,
+        fraction_train: float = 1.0,
+        fraction_evaluate: float = 1.0,
+        min_train_nodes: int = 2,
+        min_evaluate_nodes: int = 2,
+        min_available_nodes: int = 2,
+        weighted_by_key: str = "num-examples",
+        arrayrecord_key: str = "arrays",
+        configrecord_key: str = "config",
+        train_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        evaluate_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        beta: float = 0.2,
+    ) -> None:
+        super().__init__(
+            fraction_train=fraction_train,
+            fraction_evaluate=fraction_evaluate,
+            min_train_nodes=min_train_nodes,
+            min_evaluate_nodes=min_evaluate_nodes,
+            min_available_nodes=min_available_nodes,
+            weighted_by_key=weighted_by_key,
+            arrayrecord_key=arrayrecord_key,
+            configrecord_key=configrecord_key,
+            train_metrics_aggr_fn=train_metrics_aggr_fn,
+            evaluate_metrics_aggr_fn=evaluate_metrics_aggr_fn,
+        )
+        self.beta = beta
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        log(INFO, "\t├──> FedTrimmedAvg settings:")
+        log(INFO, "\t|\t└── beta: %s", self.beta)
+        super().summary()
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        # Call FedAvg aggregate_train to perform validation and aggregation
+        valid_replies, _ = self._check_and_log_replies(replies, is_train=True)
+
+        if not valid_replies:
+            return None, None
+
+        # Aggregate ArrayRecords using trimmed mean
+        # Get the key for the only ArrayRecord from the first Message
+        record_key = list(valid_replies[0].content.array_records.keys())[0]
+        # Preserve keys for arrays in ArrayRecord
+        array_keys = list(valid_replies[0].content[record_key].keys())
+
+        # Compute trimmed mean for each layer and construct ArrayRecord
+        arrays = ArrayRecord()
+        for array_key in array_keys:
+            # Get the corresponding layer from each client
+            layers = [
+                cast(ArrayRecord, msg.content[record_key]).pop(array_key).numpy()
+                for msg in valid_replies
+            ]
+            # Compute trimmed mean and save as Array in ArrayRecord
+            try:
+                arrays[array_key] = Array(trim_mean(np.stack(layers), self.beta))
+            except ValueError as e:
+                raise AggregationError(
+                    f"Trimmed mean could not be computed. "
+                    f"Likely cause: beta={self.beta} is too large."
+                ) from e
+
+        # Aggregate MetricRecords
+        metrics = self.train_metrics_aggr_fn(
+            [msg.content for msg in valid_replies],
+            self.weighted_by_key,
+        )
+        return arrays, metrics
+
+
+def trim_mean(array: NDArray, cut_fraction: float) -> NDArray:
+    """Compute trimmed mean along axis=0.
+
+    It is based on the scipy implementation:
+
+    https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.trim_mean.html
+    """
+    axis = 0
+    nobs = array.shape[0]
+    lowercut = int(cut_fraction * nobs)
+    uppercut = nobs - lowercut
+    if lowercut > uppercut:
+        raise ValueError("Fraction too big.")
+
+    atmp = np.partition(array, (lowercut, uppercut - 1), axis)
+
+    slice_list = [slice(None)] * atmp.ndim
+    slice_list[axis] = slice(lowercut, uppercut)
+    result: NDArray = np.mean(atmp[tuple(slice_list)], axis=axis)
+    return result