Mixed Precision Dynamic Loss Scale Update

Problem

Dynamic loss scaling halves the scale on overflow and doubles it after a window of consecutive finite steps without overflow.

Given the current scale, an overflow flag, the count of consecutive finite steps, and the window, return the updated scale.

Hint

Overflow takes priority over the stability window.

case 1

in: [1024,true,0,2000]

expected: 512

case 2

in: [1024,false,2000,2000]

expected: 2048

case 3

in: [512,false,500,2000]

expected: 512

Starter codePython

def mixed_precision_loss_scale_update(scale, overflow, stable_steps, window):
    # Your code here
    pass

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