Implement einsum (#2686)

* hopeful impl for Tensor.einsum

* satisfy mypy by having less typing. :(

* a few simple tests

* even more tests

* permute tests

* xfails for improper usage

* fix LLVM test fail

* use argfix

* more helpful error message on shape mismatch

test/test_ops.py

@@ -456,6 +456,56 @@ class TestOps(unittest.TestCase):
     helper_test_op([(10,20)], lambda x: x.argmin(1, False), lambda x: x.argmin(1, False), forward_only=True)
     helper_test_op([(10,20)], lambda x: x.argmin(1, True), lambda x: x.argmin(1, True), forward_only=True)
 
+  def test_einsum(self):
+    # matrix transpose
+    helper_test_op([(150,150)], lambda a: torch.einsum('ij->ji', a), lambda a: Tensor.einsum('ij->ji', a))
+    # sum all elements
+    helper_test_op([(20,30,40)], lambda a: torch.einsum('ijk->', a), lambda a: Tensor.einsum('ijk->', a))
+    # column sum
+    helper_test_op([(50,50)], lambda a: torch.einsum('ij->j', a), lambda a: Tensor.einsum('ij->j', a))
+    # row sum
+    helper_test_op([(15,15)], lambda a: torch.einsum('ij->i', a), lambda a: Tensor.einsum('ij->i', a))
+    # matrix-vector multiplication
+    helper_test_op([(15,20), (20,)], lambda a,b: torch.einsum('ik,k->i', a,b), lambda a,b: Tensor.einsum('ik,k->i', a, b))
+    # matrix-matrix multiplication
+    helper_test_op([(15,20), (20,30)], lambda a,b: torch.einsum('ik,kj->ij', a,b), lambda a,b: Tensor.einsum('ik,kj->ij', a, b))
+    # dot product
+    helper_test_op([(30),(30)], lambda a,b: torch.einsum('i,i->i', [a,b]), lambda a,b: Tensor.einsum('i,i->i', [a,b]))
+    # hadamard product
+    helper_test_op([(30,40),(30,40)], lambda a,b: torch.einsum('ij,ij->ij', a,b), lambda a,b: Tensor.einsum('ij,ij->ij', a,b))
+    # outer product
+    helper_test_op([(15,), (15,)], lambda a,b: torch.einsum('i,j->ij', a,b), lambda a,b: Tensor.einsum('i,j->ij',a,b))
+    # batch matrix multiplication
+    helper_test_op([(10,20,30),(10,30,40)], lambda a,b: torch.einsum('ijk,ikl->ijl', [a, b]), lambda a,b: Tensor.einsum('ijk,ikl->ijl', [a, b]))
+    # batch matrix multiplication, result permuted
+    helper_test_op([(10,20,25),(10,25,32)], lambda a,b: torch.einsum('ijk,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('ijk,ikl->jil', [a, b]))
+    # batch matrix multiplication, result & input permuted
+    helper_test_op([(20,10,25),(10,25,32)], lambda a,b: torch.einsum('jik,ikl->jil', [a, b]), lambda a,b: Tensor.einsum('jik,ikl->jil', [a, b]))
+    # tensor contraction
+    helper_test_op([(3,5,8,10),(11,13,5,16,8)], lambda a,b: torch.einsum('pqrs,tuqvr->pstuv', a,b), lambda a,b: Tensor.einsum('pqrs,tuqvr->pstuv', a,b), atol=1e-5)
+    # tensor contraction, input permuted
+    helper_test_op([(3,8,10,5),(11,5,13,16,8)], lambda a,b: torch.einsum('prsq,tquvr->pstuv', a,b), lambda a,b: Tensor.einsum('prsq,tquvr->pstuv', a,b), atol=1e-5)
+    # bilinear transformation
+    helper_test_op([(2,3),(5,3,7),(2,7)], lambda a,b,c: torch.einsum('ik,jkl,il->ij', [a,b,c]), lambda a,b,c: Tensor.einsum('ik,jkl,il->ij', [a,b,c]))
+
+  @unittest.expectedFailure
+  def test_einsum_shape_check(self):
+    a = Tensor.zeros(3,8,10,5)
+    b = Tensor.zeros(11,5,13,16,8)
+    Tensor.einsum('pqrs,tuqvr->pstuv',a,b)
+
+  @unittest.expectedFailure
+  def test_einsum_arity_check1(self):
+    a = Tensor.zeros(10,15)
+    b = Tensor.zeros(15,20)
+    c = Tensor.zeros(20,10)
+    Tensor.einsum('ij,jk->ij', a,b,c)
+
+  @unittest.expectedFailure
+  def test_einsum_arity_check2(self):
+    a = Tensor.zeros(10,10)
+    Tensor.einsum('ij,jk->ij', a)
+
   def test_matmul_simple(self):
     helper_test_op([(4), (4,4)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-4)
   def test_matmul(self):

tinygrad/tensor.py

@@ -511,6 +511,31 @@ class Tensor:
     return self.shape[axis]-idx.max(axis=axis, keepdim=keepdim)-1
   def argmin(self, axis=None, keepdim=False): return (-self).argmax(axis=axis, keepdim=keepdim)
 
+  @staticmethod
+  def einsum(formula, *xs) -> Tensor:
+    xs = argfix(*xs)
+    lhs, rhs = formula.split("->")
+    lhs = [sorted(enumerate(s), key=lambda e:e[1]) for s in lhs.split(',')]
+    rhs = sorted(enumerate(rhs), key=lambda e:e[1])
+
+    assert len(xs) == len(lhs), f"number of inputs doesn't match number of operands in formula, expected {len(lhs)}, got {len(xs)}"
+
+    lhs, rhs = [[list(zip(*l)) for l in lhs], list(zip(*rhs)) or [[], []]]
+    dims = {}
+    for i, l in enumerate(lhs):
+      for order, letter in enumerate(l[1]):
+        if letter not in dims: dims[letter] = xs[i].shape[l[0][order]]
+        else: assert dims[letter] == xs[i].shape[l[0][order]], f"dims of the same index should all be equal in the inputs. expected {dims[letter]} for input #{i+1}, got {xs[i].shape[l[0][order]]}"
+    xs_ = [None]*len(xs)
+    for i,x in enumerate(xs):
+      xs_[i] = x.permute(lhs[i][0]) \
+        .reshape([dims[letter] if letter in lhs[i][1] else 1 for letter in sorted(dims.keys())]) \
+        .expand([e[1] for e in sorted(dims.items(), key=lambda e: e[0])])
+
+    return reduce(lambda a,b:a*b, xs_) \
+      .sum(axis=[axis for axis,letter in enumerate(sorted(dims.keys())) if letter not in rhs[1]]) \
+      .permute(rhs[0])
+
   # ***** processing ops *****
 
   def _pool(self, k_:Tuple[sint, ...], stride:Union[Tuple[int, ...], int]=1, dilation:Union[Tuple[int, ...], int]=1) -> Tensor: