machine_learning.federated_averaging
====================================

.. py:module:: machine_learning.federated_averaging

.. autoapi-nested-parse::

   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([])  # doctest: +ELLIPSIS
   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])  # doctest: +ELLIPSIS
   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])  # doctest: +ELLIPSIS
   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]))  # doctest: +ELLIPSIS
   Traceback (most recent call last):
   ...
   ValueError: weights must be non-negative

   >>> federated_average([A, B], weights=np.array([0.0, 0.0]))  # doctest: +ELLIPSIS
   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]),
   ... )  # doctest: +ELLIPSIS
   Traceback (most recent call last):
   ...
   ValueError: weights must have shape (2,)



Functions
---------

.. autoapisummary::

   machine_learning.federated_averaging._normalize_weights
   machine_learning.federated_averaging._validate_clients
   machine_learning.federated_averaging.federated_average


Module Contents
---------------

.. py:function:: _normalize_weights(weights: numpy.ndarray, num_clients: int) -> numpy.ndarray

.. py:function:: _validate_clients(client_models: collections.abc.Sequence[collections.abc.Sequence[numpy.ndarray]]) -> None

.. py:function:: 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_models : Sequence[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.
   weights : np.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.