From 7fbebb3df6389b992f97de74b97c9d79e803ee47 Mon Sep 17 00:00:00 2001 From: Davi Silva Date: Mon, 11 Dec 2023 06:56:01 +0700 Subject: [PATCH] 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 | 50 ++++++++++++++++++++++++++++++++++++++++++++++ tinygrad/tensor.py | 25 +++++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/test/test_ops.py b/test/test_ops.py index 8f4c36097f..3c2b30ff2c 100644 --- a/test/test_ops.py +++ b/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): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 8f0cad1a77..9baf972ca5 100644 --- a/tinygrad/tensor.py +++ b/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: