from __future__ import annotations

from collections import OrderedDict

import torch


def average_weights(weight_dicts: list[OrderedDict]) -> OrderedDict:
    """Average model state dicts from multiple workers."""
    if not weight_dicts:
        raise ValueError("No weights to average")
    if len(weight_dicts) == 1:
        return weight_dicts[0]

    avg = OrderedDict()
    for key in weight_dicts[0]:
        avg[key] = torch.stack([w[key].float() for w in weight_dicts]).mean(dim=0)
    return avg