SAG Average Step

Problem

Update the running gradient sum used by the SAG method.

Replace the stale gradient old_grad with new_grad and return [new_sum, new_sum / n], where new_sum = sum_grad - old_grad + new_grad. Return a zero average when n is not positive.

Hint

Maintaining a sum avoids recomputing the average over all stored gradients.

Test cases

case 1

in: [2,0.5,1,4]

expected: [2.5,0.625]

case 2

in: [0,1,1,2]

expected: [0,0]

case 3

in: [3,0,0,0]

expected: [3,0]

Starter codePython

def sag_average_step(sum_grad, old_grad, new_grad, n):
    # Returns [new_sum, new_avg]
    # Your code here
    pass

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