Aquí está mi problema. Tengo dos matrices A y B , con entradas complejas, de dimensiones (n,n,m,m) y (n,n) respectivamente.
A continuación se muestra la operación que realizo para obtener una matriz C -
C = np.sum(B[:,:,None,None]*A, axis=(0,1)) Calcular lo anterior una vez toma alrededor de 6-8 segundos. Dado que tengo que calcular muchos de estos C , lleva mucho tiempo. ¿Hay una manera más rápida de hacer esto? (Estoy haciendo esto usando JAX NumPy en una CPU multinúcleo; NumPy normal toma incluso más tiempo)
n=77 y m=512 , si te lo estás preguntando. Puedo paralelizar mientras trabajo en un clúster, pero el gran tamaño de las matrices consume mucha memoria.
Parece que quieres einsum :
C = np.einsum('ijkl,ij->kl', A, B)Con numpy en una CPU Colab obtengo esto:
import numpy as np x = np.random.rand(50, 50, 500, 500) y = np.random.rand(50, 50) def f1(x, y): return np.sum(y[:,:,None,None]*x, axis=(0,1)) def f2(x, y): return np.einsum('ijkl,ij->kl', x, y) np.testing.assert_allclose(f1(x, y), f2(x, y)) %timeit f1(x, y) # 1 loop, best of 5: 1.52 s per loop %timeit f2(x, y) # 1 loop, best of 5: 620 ms per loop