MediumCalculusca-240

Backprop Through a Linear Layer

Problem

For a linear layer y = W x with weight matrix w (m x n), input x, and upstream gradient grad_out = dL/dy, the backward pass is dL/dW = grad_out x^T and dL/dx = W^T grad_out.

Given grad_out, w, and x, return [dW, dx] as a list containing a nested list and a flat list.

Hint

The weight gradient is the outer product of the upstream gradient and the input.

Test cases

case 1

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

expected: [[[3,4],[6,8]],[1,2]]

case 2

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

expected: [[[1,1]],[2,3]]

case 3

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

expected: [[[2,1,-1],[-4,-2,2]],[-3.5,3,-1]]

Starter codePython

def linear_layer_backward(grad_out, w, x):
    # Return [dW, dx]
    # Your code here
    pass

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