Einstein summation notation

[brandongc]

NumPy has a convenient Einstein summation capability numpy.einsum. Any interest in an API which operated in a similar way? ie

program einstein_summation
  implicit none
  integer :: i,j,k,l
  real, dimension(3,4,5) :: A = reshape([(i, i=1,60,1)], [3,4,5])
  real, dimension(4,3,2) :: B = reshape([(i, i=1,24,1)], [4,3,2])
  real, dimension(5,2) :: C = 0.0

  ! proposed API
  !C = einsum("ijk, jil-> kl", A, B)

  ! instead of 
  do k=1,5
     do l=1,2
        do i=1,3
           do j=1,4
              C(k,l) = C(k,l) + A(i,j,k) * B(j,i,l)
           end do
        end do
     end do
  end do
end program einstein_summation

Optionals:

"optimize" optional argument from NumPy's API
i.e.

C = einsum("ijk, jil-> kl", A, B, optimize=.true.)

Comments:
High performance implementations could be realized with libraries such as cuTENSOR

[ivan-pi]

Hi Brandon, I think that would be a great addition to stdlib.

Have you seen the lecture from Patrick Seewald (@pseewald) at FortranCon? (the PDF can be found here, and a recording is available on Youtube)

Example code is in the repository:
https://github.com/pseewald/fortran-einsum-example

His API uses arrays of integers for the indices, but I guess its possible to create a small parser which translates the index character string to integer indexes, and stops execution in case the user requests an invalid contraction.

A sparse tensor contraction API is part of the DBCSR (https://github.com/cp2k/dbcsr) library.

[brandongc]

Thanks for the pointer! I had not seen that talk or code.

[ivan-pi]

Just found a nice Julia package called OMEinsum.jl which can also serve as interface inspiration.

[Beliavsky]

A blog post about einsum in NumPy may give some ideas on what to implement.