如何用numexpr实现numpy.einsum三种指定表达式的等价运算？

I totally get the struggle with MemoryError when working with large arrays—numexpr is such a solid choice here since it handles chunked computation under the hood to save memory. Let’s walk through each of your einsum expressions and map them to equivalent numexpr code, with clear breakdowns of how they match up:

1. np.einsum('i, j -> ij', a, b) #

This computes the outer product of two 1D arrays: a (shape (N,)) and b (shape (M,)), resulting in a 2D array of shape (N, M).

In numexpr, you can replicate this by explicitly adding a singleton dimension to one of the arrays to trigger broadcasting (just like einsum does implicitly):

numexpr.evaluate("a[:, None] * b")

Alternatively, you could write it as:

numexpr.evaluate("a * b[None, :]")

Both work because the singleton dimension tells numexpr to stretch the array across the missing axis to match the other array’s shape.

2. np.einsum('ij, i -> ij', a, b) #

Here, you’re scaling each row of a 2D array a (shape (N, M)) by the corresponding element in a 1D array b (shape (N,)).

For numexpr, just reshape b into a column vector to align it with the rows of a:

numexpr.evaluate("a * b[:, None]")

The [:, None] turns b into shape (N, 1), which automatically broadcasts across all columns of a—exactly the same row-wise scaling as the einsum call.

3. np.einsum('ijk, kl -> ijkl', a, b) #

This takes a 3D array a (shape (I, J, K)) and a 2D array b (shape (K, L)), multiplying each slice a[i,j,:] with b to produce a 4D result of shape (I, J, K, L).

In numexpr, add a singleton dimension to the end of a to match the fourth dimension of the output:

numexpr.evaluate("a[:, :, :, None] * b")

Expanding a to shape (I, J, K, 1) lets it broadcast with b across the new fourth dimension, replicating the element-wise multiplication pattern defined by the einsum indices perfectly.

内容的提问来源于stack exchange，提问作者konstant