SWA Average Of Iterates

Problem

Average a list of parameter vectors from stochastic weight averaging.

Each iterate is a list of coordinates; return a list whose i-th entry is the arithmetic mean of coordinate i across all iterates. An empty list of iterates returns [].

Hint

Average each coordinate independently over all iterates.

Test cases

case 1

in: [[[1,2],[3,4],[5,6]]]

expected: [3,4]

case 2

in: [[[0,0],[2,4]]]

expected: [1,2]

case 3

in: [[]]

expected: []

case 4

in: [[[1],[2],[3],[4]]]

expected: [2.5]

Starter codePython

def swa_average(iterates):
    # Your code here
    pass

4 test cases · run your code against them in the editor