Trace of A Transpose A

Problem

Compute the trace of A^T A, which equals the sum of the squares of every entry of A and also the squared Frobenius norm of A.

Hint

trace(A^T A) = sum_ij A[i][j]^2.

case 1

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

expected: 30

case 2

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

expected: 2

case 3

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

expected: 14

case 4

in: [[[0]]]

expected: 0

Starter codePython

def trace_ata(A):
    # Your code here
    pass

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