英文:
Product and sum without addtional memory allocation
问题
在这个示例中，当计算C时，会创建一个大小为10x10x10x10x10的中间矩阵，然后立即折叠它。在NumPy中有没有办法避免这种情况？
英文:
Is there away to multiply two arrays and sum along an axis (or multiple axes) without allocating extra memory?
In this example:
import numpy as np
A = np.random.random((10, 10, 10))
B = np.random.random((10, 10, 10))
C = np.sum(A[:, None, :, :, None] * B[None, :, None, :, :], axis=(-1,-2))
When computing C, an intermediate matrix of size 10x10x10x10x10 is created only to be collapsed immediately. Is there a way to avoid this in numpy?
答案1
得分: 4
这看起来像是对第二个数组的转置进行的点积：
C = A @ B.T
注意：原始问题中的操作是：C = np.sum(A[:, None, :] * B[None, :, :], axis=-1)。
快速检查：
C1 = np.sum(A[:, None, :] * B[None, :, :], axis=-1)
C2 = A @ B.T
assert np.allclose(C1, C2)
您还可以使用 einsum 进行泛化：
np.einsum('ikl,jlm->ijk', A, B)
快速检查：
A = np.random.random((2, 3, 4))
B = np.random.random((2, 4, 5))
# i j k l m i j k l m
C1 = np.sum(A[:, None, :, :, None] * B[None, :, None, :, :], axis=(-1,-2))
C2 = np.einsum('ikl,jlm->ijk', A, B)
assert np.allclose(C1, C2)
英文:
This looks like a dot product with the transpose of the second array:
C = A @ B.T
NB. The operation in the original question was: C = np.sum(A[:, None, :] * B[None, :, :], axis=-1).
Quick check:
C1 = np.sum(A[:, None, :] * B[None, :, :], axis=-1)
C2 = A @ B.T
assert np.allclose(C1, C2)
You can generalize with einsum:
np.einsum('ikl,jlm->ijk', A, B)
Quick check:
A = np.random.random((2, 3, 4))
B = np.random.random((2, 4, 5))
# i j k l m i j k l m
C1 = np.sum(A[:, None, :, :, None] * B[None, :, None, :, :], axis=(-1,-2))
C2 = np.einsum('ikl,jlm->ijk', A, B)
assert np.allclose(C1, C2)
通过集体智慧和协作来改善编程学习和解决问题的方式。致力于成为全球开发者共同参与的知识库，让每个人都能够通过互相帮助和分享经验来进步。
评论