einsum elementwise addition feature request

This is a follow up to my stack overflow question. I propose adding elementwise addition and subtraction functionality to einsum. The API would remain pretty mimimal, with the just addition of the + and - characters. My gemm question would become D = np.einsum(',ij,jk+,ij', alpha, A, B, beta, C).

The parsing/API would be that on each side of a + or -, the operations must result in an array of equal shape which are then added or subtracted elementwise. A very simple un-optimized implementation could parse the above as D = np.einsum(',ij,jk', alpha, A, B) + np.einsum(',ij', beta, C)

Many optimization opportunities exist, including most of the one mentioned in this NEP. The example problem in the beginning could be A = np.einsum('...+...+...', B, C, D). It could be optimized to a single pass without intermediate arrays. It might be easier to implement those optimizations inside einsum since it is somewhat more isolated, but I don’t understand the internals nearly well enough to comment intelligently. Maybe D = np.einsum('i,ij->ij+ij', A[0,:], B, C) could use fused multiply add

Even without clever optimizations, I think the additional functionality would have significant benefits, making the most powerful linear algebra tool I know of even more comprehensive.

I am not sure if there are any issues with the API I naively suggested above, but thought it was worth suggesting. Thank you

[MichaelSnowden]

Any suggested workaround for the lack of this feature? I currently have a problem where I need to do something along the lines of either np.einsum(lambda ...) or np.einsum(..., add_to_out=x). It’s a pretty complicated question, but if I had an einsum where I could either specify the operation done or an add_to_out argument, then it would be solved.

[mattharrigan]

I don’t know how large any performance gains are or how useful they are to the community at large. I think einsum should support elementwise addition mostly for completeness. It can already support matrix and elementwise multiplication and reduction by sum. Supporting elementwise addition would complete most of the common core linear algebra operations in blas 1 through 3 with a (hopefully still) very elegant syntax. Performance gains are also nice though.

In case its helpful, here is another example of what I am proposing: A = np.einsum(‘ij,jk+ij,jk’, B, C, D, E) where the inputs are 2D arrays of compatible dimensions. I think the meaning is clear if you already know a bit about the current einsum API, A = B * C + D * E

One other point is that elementwise addition is used in the Einstein notation context. My background is originally mechanical engineering, and this notation is used in stress analysis. See this for some examples from that field.

To be clear, I do like the idea of specifying arbitrary broadcasting arithmetic using similar syntax (without the commas).

I agree also, but that would be a completely different API. My initial goal in this proposal was a narrower addition (pun intended) to an existing function.