Straight-Through Estimator Check

Problem

For the forward quantization y = round(x) * w with straight-through gradients, the backward pass treats round as the identity for x but uses the quantized value for w.

Given x, w, and the upstream gradient grad_out, return [dL/dx, dL/dw] = [grad_out * w, grad_out * round(x)].

Hint

The straight-through estimator passes the gradient through the non-differentiable rounding.

Test cases

case 1

in: [2.3,0.5,2]

expected: [1,4]

case 2

in: [2.7,3,1]

expected: [3,3]

case 3

in: [-1.4,2,1]

expected: [2,-1]

case 4

in: [0.6,1,1]

expected: [1,1]

Starter codePython

def ste_quantized_backward(x, w, grad_out):
    # Return [dL_dx, dL_dw]
    # Your code here
    pass

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