Gradient of Frobenius Norm Squared

Problem

Return the gradient of the squared Frobenius norm f(A) = sum over i, j of A[i][j]^2 with respect to A:

grad = 2 * A

This is the matrix analogue of d/dx of x^2.

Hint

Differentiate each entry independently; the result scales A by 2.

case 1

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

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

case 2

in: [[[0,0],[0,0]]]

expected: [[0,0],[0,0]]

case 3

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

expected: [[-2,0],[0,2]]

case 4

in: [[[1.5]]]

expected: [[3]]

Starter codePython

def grad_frobenius_sq(A):
    # Your code here
    pass

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