Repository: pytorch/pytorch

Issue:
Support for einsum notation
As mentioned [here](https://discuss.pytorch.org/t/einstein-summation-in-pytorch/2492), support for [einsum notation](https://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html) would be useful. It's in fact a feature that I am missing most after transitioning from TensorFlow. What's needed for the implementation is parsing the einsum string and deciding which PyTorch modules to use for which cases as well as automatically reshaping input tensor back and forth like done [here](https://github.com/tensorflow/tensorflow/blob/r1.2/tensorflow/python/ops/special_math_ops.py#L85).


Discussion:
Yeah, Pytorch lacks this. I can implement the feature
@vlasenkov that'd be great. please go ahead.
Any progress @vlasenkov ?
@MadcowD unfortunately I'm not able to contribute now, you can check the PR and complete the feature

Resolving patch:
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
--- a/tools/autograd/gen_python_functions.py
+++ b/tools/autograd/gen_python_functions.py
@@ -216,6 +216,7 @@ def create_python_bindings(python_functions, has_self, is_module=False):
         'int64_t': 'toInt64',
         'bool': 'toBool',
         'double': 'toDouble',
+        'std::string': 'string',
     }
 
     unpack_with_default_methods = {
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -1489,6 +1489,67 @@
         - **v** (*Tensor*): the eigenvectors of ``a`` if ``eigenvectors`` is ``True``; otherwise an empty tensor
 """)
 
+add_docstr(torch.einsum,
+           r"""
+einsum(equation, operands) -> Tensor
+
+This function provides a way of computing multilinear expressions (i.e. sums of products) using the
+Einstein summation convention.
+
+Args:
+    equation (string): The equation is given in terms of lower case letters (indices) to be associated
+           with each dimension of the operands and result. The left hand side lists the operands
+           dimensions, separated by commas. There should be one index letter per tensor dimension.
+           The right hand side follows after `->` and gives the indices for the output.
+           If the `->` and right hand side are omitted, it implicitly defined as the alphabetically
+           sorted list of all indices appearing exactly once in the left hand side.
+           The indices not apprearing in the output are summed over after multiplying the operands
+           entries.
+           `einsum` does not implement diagonals (multiple occurences of a single index for one tensor,
+           e.g. `ii->i`) and ellipses (`...`).
+    operands (list of Tensors): The operands to compute the Einstein sum of.
+           Note that the operands are passed as a list, not as individual arguments.
+
+Examples::
+
+    >>> x = torch.randn(5)
+    >>> y = torch.randn(4)
+    >>> torch.einsum('i,j->ij', (x,y))  # outer product
+
+    -1.0066 -2.0433 -0.8290  0.8429
+    -0.5106 -1.0365 -0.4205  0.4275
+     0.4174  0.8473  0.3438 -0.3495
+    -0.4578 -0.9292 -0.3770  0.3833
+    -0.8996 -1.8262 -0.7409  0.7533
+    [torch.FloatTensor of size (5,4)]
+
+    >>> A = torch.randn(3,5,4)
+    >>> l = torch.randn(2,5)
+    >>> r = torch.randn(2,4)
+    >>> torch.einsum('bn,anm,bm->ba', (l,A,r)) # compare torch.nn.functional.bilinear
+
+    -1.3778  2.7663 -4.9150
+    -1.7813 -4.9015  2.4149
+    [torch.FloatTensor of size (2,3)]
+
+    >>> As = torch.randn(3,2,5)
+    >>> Bs = torch.randn(3,5,4)
+    >>> torch.einsum('bij,bjk->bik', (As, Bs)) # batch matrix multiplication
+
+    (0 ,.,.) =
+     -2.0810  4.7334  2.9593  0.5268
+      1.8096 -4.6701 -2.4214 -2.2638
+
+    (1 ,.,.) =
+      0.9456 -8.3309 -2.4690 -3.3164
+      1.9580 -1.8447 -1.4268 -2.5414
+
+    (2 ,.,.) =
+     -0.1725  0.7317 -0.2110 -0.0522
+      2.5407 -0.2854  3.8720  0.9073
+    [torch.FloatTensor of size (3,2,4)]
+""")
+
 add_docstr(torch.eq,
            r"""
 eq(input, other, out=None) -> Tensor


Test patch:
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1309,6 +1309,57 @@ def test_cmul(self):
     def test_cpow(self):
         self._test_cop(torch.pow, lambda x, y: float('nan') if x < 0 else math.pow(x, y))
 
+    @unittest.skipIf(not TEST_NUMPY, 'Numpy not found')
+    def test_einsum(self):
+        # test cases taken from https://gist.github.com/rockt/15ee013889d65342088e9260a377dc8f
+        x = torch.randn(5)
+        y = torch.randn(7)
+        A = torch.randn(3, 5)
+        B = torch.randn(2, 5)
+        C = torch.randn(2, 3, 5)
+        D = torch.randn(2, 5, 7)
+        E = torch.randn(7, 9)
+        F = torch.randn(2, 3, 5, 7)
+        G = torch.randn(7, 11, 13)
+        l = torch.randn(5, 10)
+        r = torch.randn(5, 20)
+        w = torch.randn(30, 10, 20)
+        test_list = [
+            # -- Vector
+            ("i->", x),                 # sum
+            ("i,i->", x, x),            # dot
+            ("i,i->i", x, x),           # vector element-wise mul
+            ("i,j->ij", x, y),          # outer
+            # -- Matrix
+            ("ij->ji", A),              # transpose
+            ("ij->j", A),               # row sum
+            ("ij->i", A),               # col sum
+            ("ij,ij->ij", A, A),        # matrix element-wise mul
+            ("ij,j->i", A, x),          # matrix vector multiplication
+            ("ij,kj->ik", A, B),        # matmul
+            ("ij,ab->ijab", A, E),      # matrix outer product
+            # -- Tensor
+            ("aij,ajk->aik", C, D),     # batch matmul
+            ("ijk,jk->i", C, A),        # tensor matrix contraction
+            ("aij,jk->aik", D, E),      # tensor matrix contraction
+            ("abcd,dfg->abcfg", F, G),  # tensor tensor contraction
+            ("ijk,jk->ik", C, A),       # tensor matrix contraction with double indices
+            ("ijk,jk->ij", C, A),       # tensor matrix contraction with double indices
+            ("ijk,ik->j", C, B),        # non contiguous
+            ("ijk,ik->jk", C, B),       # non contiguous with double indices
+            # -- Other
+            ("bn,anm,bm->ba", l, w, r),  # as torch.bilinear
+        ]
+        for test in test_list:
+            actual = torch.einsum(test[0], test[1:])
+            expected = np.einsum(test[0], *[t.numpy() for t in test[1:]])
+            self.assertEqual(expected.shape, actual.shape)
+            self.assertTrue(np.allclose(expected, actual.numpy()))
+
+            def do_einsum(*args):
+                return torch.einsum(test[0], args)
+            self.assertTrue(torch.autograd.gradcheck(do_einsum, test[1:]))
+
     def test_sum_all(self):
         def check_sum_all(tensor):
             pylist = tensor.reshape(-1).tolist()