Vectorized computation on Numpy array

Currently, I am using the loops to do the task below. It is quite similar to my previous question: Dot product two 4D Numpy array

Just wondering if anyone can help me to vectorize the computation without the loops.

A1 = np.random.rand(2, 320, 320)
A2 = np.random.rand(15, 2, 320, 320)
B   = np.zeros((2, 320, 320))

for row in range (320):
    for col in range (320):
        C = np.diag(A1[:, row, col]) - (np.dot(A2[:, :, row, col].T, A2[:, :, row, col]) / 15) # the shape of array C is (2, 2)
        C = np.diag(C)                
        B[:, row, col] = C

Any help or suggestions would be greatly appreciated!