a=np.array([[1,2,3],[4,5,6]])

np.square(a) 
array([[ 1,  4,  9],
       [16, 25, 36]])
       
np.einsum('ij,ij->ij',a,a)
Out: array([[ 1,  4,  9],
       [16, 25, 36]])     

np.diag(np.dot(a, (a).T))
Out: array([14, 77])

np.einsum('ij,ij->i',a,a)
Out: array([14, 77])

np.sum(np.square(a))
Out: 91

np.einsum('ij,ij',a,a)
Out: 91

链接： https://zhuanlan.zhihu.com/p/27739282