为了说明我面临的问题，下面是一些示例代码:
a = np.round(np.random.rand(10, 15))
counta = np.count_nonzero(a, axis=-1)
print(counta)
A = np.einsum('im,mj->ijm', a, a.T)
countA = np.count_nonzero(A, axis=-1)
print(countA)
它创建一个二维数组，并沿最后一个轴计算其非零元素.然后，它创建一个3D数组，其中的非零元素沿最后一个轴再次计数.
我的问题是，我的数组a太大了，我可以执行第一步，但不能执行第二步，因为A数组会占用太多内存.
Is there any way to still get 100? That is to count the zeros in A along a given axis without actually creating the array?
我认为您可以简单地使用矩阵乘法(点积)来获得结果，而不需要生成庞大的3D数组A:
a = np.round(np.random.rand(10, 15)).astype(int)
counta = np.count_nonzero(a, axis=-1)
A = np.einsum('im,mj->ijm', a, a.T)
countA = np.count_nonzero(A, axis=-1)
assert np.all(countA == (a @ a.T))
这也要快得多:
a = np.round(np.random.rand(1000, 1500)).astype(int)
%timeit np.count_nonzero(np.einsum('im,mj->ijm', a, a.T), axis=-1)
3.94 s ± 38.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit a @ a.T
558 ms ± 6.74 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
还要注意，第一步是多余的，第二步是多余的:
assert np.all(counta == np.diag(countA))