From 55707fd00de47388cb1f15322d802f65fc50fbc7 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 17 Nov 2024 10:58:41 -0500 Subject: [PATCH] fix passing sum_acc_dtype="" to Tensor.sum should fail (#7748) --- test/test_ops.py | 6 ++++++ tinygrad/tensor.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/test/test_ops.py b/test/test_ops.py index 45359c389c..4d3b99e987 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -1011,6 +1011,12 @@ class TestOps(unittest.TestCase): self.helper_test_exception([()], lambda x: x.sum(1), lambda x: x.sum(1), expected=IndexError) self.helper_test_exception([()], lambda x: x.sum((1,)), lambda x: x.sum((1,)), expected=IndexError) + def test_sum_acc_dtype(self): + helper_test_op([(45,3)], lambda x: x.sum(), lambda x: x.sum(acc_dtype=dtypes.float32)) + if is_dtype_supported(dtypes.float64): helper_test_op([(45,3)], lambda x: x.sum(dtype=torch.float64), lambda x: x.sum(acc_dtype=dtypes.float64)) + + with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).sum(acc_dtype="") + def test_sum_with_zeros_shape(self): helper_test_op([(4, 0)], lambda x: x.sum(axis=(0,))) helper_test_op([(4, 0)], lambda x: x.sum(axis=(1,))) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index a5808b9a08..ea51eca86c 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1526,7 +1526,7 @@ class Tensor(SimpleMathTrait): # pylint: disable=abstract-method print(t.sum(axis=1).numpy()) ``` """ - ret = self.cast(acc_dtype or sum_acc_dtype(self.dtype))._reduce(F.Sum, axis, keepdim) + ret = self.cast(sum_acc_dtype(self.dtype) if acc_dtype is None else acc_dtype)._reduce(F.Sum, axis, keepdim) return ret.cast(self.dtype) if acc_dtype is None and self.dtype in (dtypes.float16, dtypes.bfloat16) else ret def prod(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False, acc_dtype:Optional[DTypeLike]=None):