diff --git a/test/test_ops.py b/test/test_ops.py index 192788cdb7..0350b635aa 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -2559,12 +2559,16 @@ class TestOps(unittest.TestCase): lambda x,src: x.scatter(dim=1, index=a, src=src), forward_only=True) helper_test_op([(10,3,10), (10,10,10)], lambda x,src: x.scatter(dim=1, index=b, src=src), lambda x,src: x.scatter(dim=1, index=a, src=src), forward_only=True) + self.helper_test_exception([(2,3,10), (10,10,10)], lambda x,src: x.scatter(dim=1, index=b, src=src), lambda x,src: x.scatter(dim=1, index=a, src=src), expected=(RuntimeError, AssertionError)) self.helper_test_exception([(10,3,10), (10,3,10)], lambda x,src: x.scatter(dim=1, index=b, src=src), lambda x,src: x.scatter(dim=1, index=a, src=src), expected=(RuntimeError, AssertionError)) self.helper_test_exception([(3,4,5), (3,4,5)], lambda x,src: x.scatter(dim=1, index=b, src=src, mode="typo"), - lambda x,src: x.scatter(dim=1, index=a, src=src, mode="typo"), expected=TypeError) + lambda x,src: x.scatter(dim=1, index=a, src=src, mode="typo"), expected=TypeError) + self.helper_test_exception([(3,4,5), (3,4,5)], lambda x,src: x.half().scatter(dim=1, index=b, src=src), + lambda x,src: x.half().scatter(dim=1, index=a, src=src), expected=RuntimeError) + helper_test_op([(4,5,6)], lambda x: x.scatter(dim=1, index=b, value=3), lambda x: x.scatter(dim=1, index=a, src=3), forward_only=True) helper_test_op([(4,5,6)], lambda x: x.scatter(dim=1, index=b, value=float("inf")), lambda x: x.scatter(dim=1, index=a, src=float("inf")), forward_only=True) @@ -2626,13 +2630,19 @@ class TestOps(unittest.TestCase): lambda src: y.scatter_reduce(dim=1, index=b, src=src, reduce="prod"), lambda src: x.scatter_reduce(dim=1, index=a, src=src, reduce="prod"), forward_only=True) - def test_scatter_reduce_invalid_reduce_op(self): + def test_scatter_reduce_errors(self): b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False) a = Tensor(b.detach().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False) + # invalid reduce arg self.helper_test_exception([(4,5,6), (4,5,6)], lambda x,src: x.scatter_reduce(dim=0, index=b, src=src, reduce="INVALID"), lambda x,src: x.scatter_reduce(dim=0, index=a, src=src, reduce="INVALID"), RuntimeError) + # dtype mismatch + self.helper_test_exception([(4,5,6), (4,5,6)], + lambda x,src: x.half().scatter_reduce(dim=0, index=b, src=src, reduce="sum"), + lambda x,src: x.half().scatter_reduce(dim=0, index=a, src=src, reduce="sum"), + RuntimeError) def test_scaled_dot_product_attention(self): helper_test_op([(32,8,16,64), (32,8,16,64), (32,8,16,64)], torch.nn.functional.scaled_dot_product_attention, Tensor.scaled_dot_product_attention) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index c0e3f869cc..903e253c8c 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -2439,9 +2439,11 @@ class Tensor(SimpleMathTrait): return x.cast(self.dtype) def _pre_scatter(self, dim:int, index:Tensor, src:Tensor) -> tuple[Tensor, Tensor]: + index, dim = index.to(self.device), self._resolve_dim(dim) assert index.ndim == self.ndim == src.ndim, f"self.ndim, index.ndim and src.dim must all equal, {self.ndim=} {index.ndim=} {src.ndim=}" assert all((d == dim or self_ >= index_) and src_ >= index_ for d,(self_,index_,src_) in enumerate(zip(self.shape, index.shape, src.shape))), \ f"All dimensions of {index.shape=} should be <= to all dimensions of {src.shape=} and all dimensions except dimension {dim} of {self.shape=}" + if self.dtype != src.dtype: raise RuntimeError(f"expect {self.dtype=} to be equal to {src.dtype=}") # shrink src to index shape to shrink away the unused values src = src.shrink(tuple((0,s) for s in index.shape)) # prepare src and mask for reduce with respect to dim @@ -2479,8 +2481,7 @@ class Tensor(SimpleMathTrait): """ if reduce not in {None, "add", "multiply"}: raise TypeError(f"{reduce=} must be one of None, 'multiply', or 'add'") if reduce and isinstance(src, Tensor): raise TypeError("Tensor src is not supported with reduce arg. see scatter_reduce") - src = src.cast(self.dtype) if isinstance(src, Tensor) else Tensor(src, device=self.device, dtype=self.dtype)._broadcast_to(index.shape) - index, dim = index.to(self.device), self._resolve_dim(dim) + if not isinstance(src, Tensor): src = Tensor(src, device=self.device, dtype=self.dtype)._broadcast_to(index.shape) src, mask = self._pre_scatter(dim, index, src) # TODO: should not overwrite acc_dtype here? if reduce == "add": return mask.where(src, 0).sum(-1, acc_dtype=self.dtype) + self @@ -2517,8 +2518,6 @@ class Tensor(SimpleMathTrait): print(Tensor([[-10, 20, 0, 5, 10]], dtype=src.dtype).scatter_reduce(0, index, src, reduce='amin').numpy()) ``` """ - src = src.cast(self.dtype) - index, dim = index.to(self.device), self._resolve_dim(dim) src, mask = self._pre_scatter(dim, index, src) def _inv_mask(a:Union[Tensor, ConstType], b:Union[Tensor, ConstType]) -> Tensor: return mask.any(-1).logical_not().where(a, b) # TODO: should not overwrite acc_dtype here?