Straight-Through Estimator Gradient

Problem

Apply the straight-through estimator (STE) gradient used in quantization-aware training: pass grad_output through unchanged where |x| <= clip, and return 0 outside the clipping range. Inputs are element-wise lists and the result is a list of the same length.

The signature is ste_gradient(x, grad_output, clip=1.0).

Hint

The rounding function has zero derivative almost everywhere, so STE pretends it is identity.

Test cases

case 1

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

expected: [1,2,0,0]

case 2

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

expected: [1,2,0]

case 3

in: [[0,0.25,0.5],[5,6,7],1]

expected: [5,6,7]

case 4

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

expected: [0]

Starter codePython

def ste_gradient(x, grad_output, clip=1.0):
    # Your code here
    pass

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