Mixed Precision Loss Scale Update

Problem

Update a mixed-precision loss scale for one optimizer step. If any gradient overflowed (found_inf truthy), multiply the loss scale by the backoff factor; otherwise multiply it by the growth factor.

The signature is mixed_precision_loss_scale(loss_scale, found_inf, growth_factor=2.0, backoff_factor=0.5). Return the new loss scale.

Hint

Overflow backs the scale off; a clean step lets it grow again.

Test cases

case 1

in: [1024,1]

expected: 512

case 2

in: [1024,0]

expected: 2048

case 3

in: [65536,0,1.5]

expected: 98304

case 4

in: [100,1,2,0.1]

expected: 10

Starter codePython

def mixed_precision_loss_scale(loss_scale, found_inf, growth_factor=2.0, backoff_factor=0.5):
    # Your code here
    pass

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