Gradient of trace(AX)

Problem

For the scalar function f(X) = trace(A*X), return its gradient with respect to X:

grad = A^T

A and X have compatible square shapes.

Hint

trace(AX) = sum_ij A[i][j] X[j][i], so the derivative with respect to X[j][i] is A[i][j].

case 1

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

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

case 2

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

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

case 3

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

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

case 4

in: [[[0]]]

expected: [[0]]

Starter codePython

def grad_trace_ax(A):
    # Your code here
    pass

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