了解NumPy的einsum [英] Understanding NumPy's einsum

问题描述

我正在努力确切地了解einsum的工作方式.我看了一下文档和一些示例，但看起来似乎不合时宜.

I'm struggling to understand exactly how einsum works. I've looked at the documentation and a few examples, but it's not seeming to stick.

这是我们在课堂上看过的一个例子:

Here's an example we went over in class:

C = np.einsum("ij,jk->ki", A, B)

用于两个数组A和B

我认为这会花费A^T * B，但是我不确定(它正在对其中之一进行转置吗?).谁能告诉我这里到底发生了什么(以及通常使用einsum时)?

I think this would take A^T * B, but I'm not sure (it's taking the transpose of one of them right?). Can anyone walk me through exactly what's happening here (and in general when using einsum)?

推荐答案

(注意:此答案基于简短的关于我之前写过的einsum的博客帖子.

(Note: this answer is based on a short blog post about einsum I wrote a while ago.)

想象一下，我们有两个多维数组A和B.现在，假设我们要...

Imagine that we have two multi-dimensional arrays, A and B. Now let's suppose we want to...

将
• A与B进行特殊组合，以创建新的产品系列；然后也许
• sum 沿着特定轴的这个新数组；然后也许
• 以特定顺序转置新数组的轴.

• multiply A with B in a particular way to create new array of products; and then maybe
• sum this new array along particular axes; and then maybe
• transpose the axes of the new array in a particular order.

einsum有很大的机会帮助我们更快，更有效地执行此操作，如multiply，sum和transpose的NumPy函数的组合将允许.

There's a good chance that einsum will help us do this faster and more memory-efficiently that combinations of the NumPy functions like multiply, sum and transpose will allow.

这是一个简单(但并非完全无关紧要)的例子.取以下两个数组:

Here's a simple (but not completely trivial) example. Take the following two arrays:

A = np.array([0, 1, 2])

B = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])

我们将A和B逐个元素相乘，然后沿着新数组的行求和.在正常"的NumPy中，我们将编写:

We will multiply A and B element-wise and then sum along the rows of the new array. In "normal" NumPy we'd write:

>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])

因此，在这里，在A上的索引操作将两个数组的第一个轴对齐，以便可以广播乘法.然后将产品数组中的行相加，以得出答案.

So here, the indexing operation on A lines up the first axes of the two arrays so that the multiplication can be broadcast. The rows of the array of products is then summed to return the answer.

现在，如果我们想改用einsum，我们可以这样写:

Now if we wanted to use einsum instead, we could write:

>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])

签名字符串'i,ij->i'是此处的关键，需要一些解释.您可以将其分为两半.在左侧(->的左侧)，我们标记了两个输入数组.在->的右侧，我们标记了要结束的数组.

The signature string 'i,ij->i' is the key here and needs a little bit of explaining. You can think of it in two halves. On the left-hand side (left of the ->) we've labelled the two input arrays. To the right of ->, we've labelled the array we want to end up with.

这是接下来发生的事情:

Here is what happens next:

• A具有一个轴；我们已将其标记为i. B有两个轴；我们已将轴0标记为i，将轴1标记为j.

• A has one axis; we've labelled it i. And B has two axes; we've labelled axis 0 as i and axis 1 as j.

在两个输入数组中通过重复标签i，我们告诉einsum这两个轴应相乘在一起.换句话说，我们将数组A与数组B的每一列相乘，就像A[:, np.newaxis] * B一样.

By repeating the label i in both input arrays, we are telling einsum that these two axes should be multiplied together. In other words, we're multiplying array A with each column of array B, just like A[:, np.newaxis] * B does.

请注意，j不会在所需的输出中显示为标签；我们只使用了i(我们想以一维数组结尾).通过省略标签，我们告诉einsum沿该轴求和.换句话说，就像.sum(axis=1)一样，我们正在对产品的行进行求和.

Notice that j does not appear as a label in our desired output; we've just used i (we want to end up with a 1D array). By omitting the label, we're telling einsum to sum along this axis. In other words, we're summing the rows of the products, just like .sum(axis=1) does.

使用einsum基本上就是您所需要知道的.玩一会会有所帮助；如果我们将两个标签都留在输出'i,ij->ij'中，则会得到二维的产品数组(与A[:, np.newaxis] * B相同).如果我们说没有输出标签'i,ij->，我们将得到一个数字(与(A[:, np.newaxis] * B).sum()相同).

That's basically all you need to know to use einsum. It helps to play about a little; if we leave both labels in the output, 'i,ij->ij', we get back a 2D array of products (same as A[:, np.newaxis] * B). If we say no output labels, 'i,ij->, we get back a single number (same as doing (A[:, np.newaxis] * B).sum()).

einsum的伟大之处在于，它不会首先构建临时产品数组；它只是对产品进行累加.这样可以节省大量内存.

The great thing about einsum however, is that is does not build a temporary array of products first; it just sums the products as it goes. This can lead to big savings in memory use.

为解释点积，这里有两个新数组:

To explain the dot product, here are two new arrays:

A = array([[1, 1, 1],
           [2, 2, 2],
           [5, 5, 5]])

B = array([[0, 1, 0],
           [1, 1, 0],
           [1, 1, 1]])

我们将使用np.einsum('ij,jk->ik', A, B)计算点积.这是一张显示A和B的标签以及我们从该函数获得的输出数组的图片:

We will compute the dot product using np.einsum('ij,jk->ik', A, B). Here's a picture showing the labelling of the A and B and the output array that we get from the function:

您会看到标签j被重复-这意味着我们正在将A的行与B的列相乘.此外，输出中不包含标签j-我们对这些乘积求和.标签i和k保留用于输出，因此我们获得了2D数组.

You can see that label j is repeated - this means we're multiplying the rows of A with the columns of B. Furthermore, the label j is not included in the output - we're summing these products. Labels i and k are kept for the output, so we get back a 2D array.

将结果与未加标签j的数组进行比较可能更加清楚.在下面的左侧，您可以看到写入np.einsum('ij,jk->ijk', A, B)产生的3D数组(即，我们保留了标签j):

It might be even clearer to compare this result with the array where the label j is not summed. Below, on the left you can see the 3D array that results from writing np.einsum('ij,jk->ijk', A, B) (i.e. we've kept label j):

求和轴j给出了预期的点积，如右图所示.

Summing axis j gives the expected dot product, shown on the right.

要获得更多有关einsum的感觉，使用下标符号实现熟悉的NumPy数组操作可能会很有用.涉及乘和求和轴组合的任何内容都可以使用einsum编写.

To get more of feel for einsum, it can be useful to implement familiar NumPy array operations using the subscript notation. Anything that involves combinations of multiplying and summing axes can be written using einsum.

让A和B是两个长度相同的一维数组.例如，A = np.arange(10)和B = np.arange(5, 15).

Let A and B be two 1D arrays with the same length. For example, A = np.arange(10) and B = np.arange(5, 15).

• A的总和可以写为:

np.einsum('i->', A)

元素级乘法A * B可以写为:

Element-wise multiplication, A * B, can be written:

np.einsum('i,i->i', A, B)

内部乘积或点乘np.inner(A, B)或np.dot(A, B)可以写为:

The inner product or dot product, np.inner(A, B) or np.dot(A, B), can be written:

np.einsum('i,i->', A, B) # or just use 'i,i'

外部产品np.outer(A, B)可以写为:

The outer product, np.outer(A, B), can be written:

np.einsum('i,j->ij', A, B)

对于2D数组，C和D，只要轴是兼容的长度(相同长度或其中一个长度为1)，下面是一些示例:

For 2D arrays, C and D, provided that the axes are compatible lengths (both the same length or one of them of has length 1), here are a few examples:

• C(主对角线总和)的轨迹np.trace(C)可以写为:

• The trace of C (sum of main diagonal), np.trace(C), can be written:

np.einsum('ii', C)

C的元素级乘法和D的转置，C * D.T，可以写成:

Element-wise multiplication of C and the transpose of D, C * D.T, can be written:

np.einsum('ij,ji->ij', C, D)

将C的每个元素乘以D数组(以创建4D数组)C[:, :, None, None] * D:

Multiplying each element of C by the array D (to make a 4D array), C[:, :, None, None] * D, can be written:

np.einsum('ij,kl->ijkl', C, D)  

这篇关于了解NumPy的einsum的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！