cp2k / dbcsr

DBCSR: Distributed Block Compressed Sparse Row matrix library
https://cp2k.github.io/dbcsr/
GNU General Public License v2.0
134 stars 45 forks source link

Add einsum() API for tensor contractions #192

Open oschuett opened 5 years ago

oschuett commented 5 years ago

Machine learning libraries have found a nice API for tensor contractions, which is based on the Einstein notation. Examples are PyTorch, Tensorflow, and NumPy. I guess, it would be rather straightforward to add this to DBCSR as well.

pseewald commented 5 years ago

This would indeed facilitate the use of DBCSR tensors but there are a few things to consider:

1) Not all tensor operations that can be expressed with Einstein notation are supported. This is due to the restriction of the low-level implementations that all operations must map to either matrix-matrix multiplication and/or data redistribution. Thus the following operations are supported:

2) The notation can not be fully abstract since, for good performance, the user needs to control the matrix representation of each tensor (at creation). However that could be expressed in a similar notation: ijk <-> ki,j meaning that a tensor ijk is mapped to a matrix with first index ki and second index j.