geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

feat: add acc_dtype to einsum (#4571)

Browse Source
This commit is contained in:
Filip Brzek
2024-05-13 14:02:07 -04:00
committed by GitHub
parent d97d5a7689
commit f7d08bd454
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_linearizer.py
+1
View File
@@ -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)
tinygrad/tensor.py
+2 -2
View File
@@ -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 *****