machine_learning.federated_averaging

Federated averaging (FedAvg) utilities.

This module provides a simple NumPy-based implementation of the FedAvg aggregation algorithm. It supports equal weighting and custom non-negative weights that are normalized internally.

Doctests

Basic equal-weight averaging across two “clients” with two tensors each (vector and 2x2 matrix):

>>> A = [
...     np.array([1.0, 2.0]),
...     np.array([[1.0, 2.0], [3.0, 4.0]]),
... ]
>>> B = [
...     np.array([3.0, 4.0]),
...     np.array([[5.0, 6.0], [7.0, 8.0]]),
... ]
>>> eq = federated_average([A, B])
>>> eq[0].tolist()
[2.0, 3.0]
>>> eq[1].tolist()
[[3.0, 4.0], [5.0, 6.0]]

Weighted averaging with weights [2, 1] (normalized to [2/3, 1/3]):

>>> w = federated_average(
...     [A, B],
...     weights=np.array([2.0, 1.0]),
... )
>>> w[0].tolist()
[1.6666666666666665, 2.6666666666666665]
>>> w[1].tolist()
[[2.333333333333333, 3.333333333333333], [4.333333333333333, 5.333333333333333]]

Error cases:

  • No clients

>>> federated_average([])
Traceback (most recent call last):
...
ValueError: client_models must be a non-empty list

  • Mismatched number of tensors per client

>>> C = [np.array([1.0, 2.0])]  # only one tensor
>>> federated_average([A, C])
Traceback (most recent call last):
...
ValueError: All clients must have the same number of tensors

  • Mismatched tensor shapes across clients

>>> C2 = [
...     np.array([1.0, 2.0]),
...     np.array([[1.0, 2.0]]),
... ]  # second tensor has different shape
>>> federated_average([A, C2])
Traceback (most recent call last):
...
ValueError: Client 2 tensor shape (1, 2) does not match (2, 2)

  • Invalid weights: negative or wrong shape or zero-sum

>>> federated_average([A, B], weights=np.array([1.0, -1.0]))
Traceback (most recent call last):
...
ValueError: weights must be non-negative

>>> federated_average([A, B], weights=np.array([0.0, 0.0]))
Traceback (most recent call last):
...
ValueError: weights must sum to a positive value

>>> federated_average(
...     [A, B],
...     weights=np.array([1.0, 2.0, 3.0]),
... )
Traceback (most recent call last):
...
ValueError: weights must have shape (2,)

Functions

_normalize_weights(→ numpy.ndarray)

_validate_clients(→ None)

federated_average(→ list[numpy.ndarray])

Compute the weighted average of clients' model tensors.

Module Contents

machine_learning.federated_averaging._normalize_weights(weights: numpy.ndarray, num_clients: int) → numpy.ndarray
machine_learning.federated_averaging._validate_clients(client_models: collections.abc.Sequence[collections.abc.Sequence[numpy.ndarray]]) → None
machine_learning.federated_averaging.federated_average(client_models: collections.abc.Sequence[collections.abc.Sequence[numpy.ndarray]], weights: numpy.ndarray | None = None) → list[numpy.ndarray]

Compute the weighted average of clients’ model tensors.

Parameters

client_modelsSequence[Sequence[np.ndarray]]

A list of clients, each being a sequence of NumPy arrays (tensors). All clients must have the same number of tensors with identical shapes.

weightsnp.ndarray | None, optional

A 1-D array of non-negative weights, one per client. If None, equal weighting is used. Weights are normalized to sum to 1.

Returns

list[np.ndarray]

The list of aggregated tensors with the same shapes as the inputs.