R2D2 Burn-In Masked Loss

Problem

Average only the loss terms that come after an R2D2 burn-in prefix.

Return the mean of losses[burn_in:] with negative burn_in treated as 0. When the burn-in consumes the whole sequence or the list is empty, return 0.0.

Hint

Burn-in steps warm up the recurrent state but must not contribute gradients.

Test cases

case 1

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

expected: 3.5

case 2

in: [[5],0]

expected: 5

case 3

in: [[1,2],2]

expected: 0

case 4

in: [[],0]

expected: 0

case 5

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

expected: 2

Starter codePython

def r2d2_masked_loss(losses, burn_in):
    # Your code here
    pass

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