Python, numpy: Frobineus norm of a matrix using np.diag

I am practising with numpy and now I need to calculate the F norm of a matrix A following 2 different equations.

F_n = sqrt (sum(a^2)) (a is each element of A) AND F_n = sqrt (trace (A * At)) At is A transpose

The first one is easy


The second one is more challenging because (A*At) might be too big.

I am trying to use np.diag to avoid multiplying 2 (big) matrices. I would convert A to a vector A1 that is "A in line-by-line"; then convert A1 to A2 that A2 diagonal = A1, the rest of A2 are 0. Not sure whether it is a good move.

Anyway, in such case A2 will have a bigger size but A2t will be equal to A2, and (A2*A2t) can be easily calculated.

My attempt:

A1 = A.reshape(1,-1)
A2 = np.diag(A)
F_n = np.sum(A2*A2.T)

It does not work well. Could you please help me to improve the code.