FP8 Amax Update

Problem

Update a running FP8 amax (maximum absolute value) with a new tensor and a decay factor: max(max(|x|) over the tensor, decay * current_amax). An empty tensor leaves decay * current_amax.

The signature is fp8_amax_update(current_amax, tensor, decay=0.0).

Hint

Decay reduces the influence of stale amax history so the scale can shrink.

Test cases

case 1

in: [5,[-2,7.5,3],0]

expected: 7.5

case 2

in: [10,[1,-2],0.5]

expected: 5

case 3

in: [0,[],0.9]

expected: 0

case 4

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

expected: 4

case 5

in: [100,[],1]

expected: 100

Starter codePython

def fp8_amax_update(current_amax, tensor, decay=0.0):
    # Your code here
    pass

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