2

NumPy's einsum lets you explicitly choose which axes are contracted with the so-called explicit mode, using ->:

>>> a = np.arange(9).reshape((3, 3))
>>> np.einsum('ij,ij->j', a, a)
array([45, 66, 93])

einsum also lets you specify a programmatic interface where the axes are specified using a list of integers rather than a string.

>>> np.einsum(a, [0, 1], a, [0, 1])
204

But it isn't clear if it's possible to combine these two paradigms. I'm trying to write a function that contracts a specified axis of two arrays, where the arrays can have any dimensionality.

asmeurer
  • 86,894
  • 26
  • 169
  • 240
  • 1
    According to the docstring, "`einsum` also provides an alternative way to provide the subscripts and operands as ``einsum(op0, sublist0, op1, sublist1, ..., [sublistout])``". Does that optional final argument do what you want? E.g. `np.einsum(a, [0, 1], a, [0, 1], [1])` – Warren Weckesser Jun 30 '22 at 22:29
  • Note that you can use a transposition to do a similar thing. Transpositions create views so the content is not copied. Transposed views can be taken into account by BLAS implementation during the optimization step (if the einsum call is not optimized then it will be pretty inefficient whatever the method used) – Jérôme Richard Jun 30 '22 at 22:44
  • The list format should be able to express the same functionality as the string version. – hpaulj Jun 30 '22 at 22:50

0 Answers0