I am looking for a faster alternative for the following code.
rows = 30
cols = 10**3
v_mat = np.random.rand(rows,cols)
a_mat = np.random.rand(cols,cols)
b_vec = (v_mat @ a_mat @ v_mat.T).diagonal()
Asked
Active
Viewed 39 times
1
Chaitanya Kesanapalli
- 163
- 1
- 8
-
I edited the question. – Chaitanya Kesanapalli Sep 20 '21 at 22:05
-
@Lagerbaer Yes, it works. The reference solution is 4 times faster in my case. Thank you. – Chaitanya Kesanapalli Sep 20 '21 at 22:17
-
1`@` can be used without the diagional: `(v_mat[:,None,:]@a_mat@v_mat[:,:,None])[:,0,0]`. But `einsum` with `optimize=True` is better, as is the linked `(A.dot(B) * A).sum(-1)` – hpaulj Sep 20 '21 at 23:54
-
@hpaulj, Thanks. I will try the `einsum` code. – Chaitanya Kesanapalli Sep 21 '21 at 00:17