From f7d08bd454f5c644f3baed0731f0dbfabb93bdaf Mon Sep 17 00:00:00 2001 From: Filip Brzek Date: Mon, 13 May 2024 20:02:07 +0200 Subject: [PATCH] feat: add acc_dtype to einsum (#4571) --- test/test_linearizer.py | 1 + tinygrad/tensor.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/test/test_linearizer.py b/test/test_linearizer.py index bc4d84be8c..b76967f0a3 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -256,6 +256,7 @@ class TestLinearizer(unittest.TestCase): a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype) helper_arg_acc_dtype(a.sum(acc_dtype=acc_dtype), expected_dtype) helper_arg_acc_dtype(a.matmul(b, acc_dtype=acc_dtype), expected_dtype) + helper_arg_acc_dtype(Tensor.einsum("ki,ij->kj", a, b, acc_dtype=acc_dtype), expected_dtype) d, w = Tensor.rand(4, 8, 8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, 2, 2, dtype=tensor_dtype) helper_arg_acc_dtype(d.conv2d(w, acc_dtype=acc_dtype), expected_dtype) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index b148e938be..e897621b34 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -968,7 +968,7 @@ class Tensor: def argmin(self, axis=None, keepdim=False): return (-self).argmax(axis=axis, keepdim=keepdim) @staticmethod - def einsum(formula:str, *raw_xs) -> Tensor: + def einsum(formula:str, *raw_xs, acc_dtype:Optional[DType]=None) -> Tensor: xs:Tuple[Tensor] = argfix(*raw_xs) formula = formula.replace(" ", "") inputs_str, output = formula.split("->") if "->" in formula else (formula, sorted(formula)) @@ -990,7 +990,7 @@ class Tensor: # sum over all axes that's not in the output, then permute to the output order return functools.reduce(lambda a,b:a*b, xs_) \ - .sum(axis=[axis for axis,(letter,_) in enumerate(letter_val) if letter not in output]).permute(rhs_order) + .sum(axis=[axis for axis,(letter,_) in enumerate(letter_val) if letter not in output],acc_dtype=acc_dtype).permute(rhs_order) # ***** processing ops *****