diff --git a/leanpass/tensor.py b/leanpass/tensor.py index 6ea1cfc..64c8c35 100644 --- a/leanpass/tensor.py +++ b/leanpass/tensor.py @@ -108,6 +108,8 @@ def __rmul__(self, other): def __truediv__(self, other): other = other if isinstance(other, Tensor) else Tensor(other) + if np.any(other.data == 0): + raise ZeroDivisionError("Tensor division by zero") out = self._create_child(self.data / other.data, "/", (self, other)) def _backward(): diff --git a/tests/test_leanpass.py b/tests/test_leanpass.py index 2890999..f70bef2 100644 --- a/tests/test_leanpass.py +++ b/tests/test_leanpass.py @@ -97,3 +97,11 @@ def test_tensor_gelu_backward(): assert x.grad.shape == x.data.shape assert np.all(np.isfinite(x.grad)) assert np.all(x.grad != 0) + + +def test_truediv_raises_on_zero_divisor(): + import pytest + x = Tensor([1.0, 2.0], requires_grad=True) + y = Tensor([1.0, 0.0], requires_grad=False) + with pytest.raises(ZeroDivisionError): + _ = x / y