Быстрый способ вычислить среднее значение для веса нейронной сети

Question

Issue

Я пытаюсь вычислить среднее значение n весов нейронной сети (keras \ tenorflow). Я хочу максимально оптимизировать вычисления.

Мое решение

def get_avg_weights(array_weights):
    new_weights = list()
    for weights_list_tuple in zip(*array_weights):
        new_weights.append(
            [sum(weights_) / len(weights_) \
             for weights_ in zip(*weights_list_tuple)])
    return new_weights

if __name__ == 'main':
    array_weights = deque([], maxlen = 5)   # collection of different NN weights
    array_weights.append(weights1)
    array_weights.append(weights2)
    array_weights.append(weights3)

    avg = get_avg_weights(array_weights)

Результат правильный, но относительно медленный.
Для 5 наборов весов и NN из 4 слоев и 1000 скрытых нейронов на слой требуется 34ms для вычисления среднего. Поскольку мне нужно сделать это около 600 тыс. Раз с различным набором весов, я хочу максимально оптимизировать вычисления. Кто-нибудь знает способ ускорить код?