sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 6 Activity

fix elu/celu/selu gradient being nan on large inputs (#17852)

Browse Source 
* fix elu/celu/selu gradient being nan on large inputs

* shrink shape for the elu/celu/selu overflow checks

* route selu and celu through elu
This commit is contained in:
Teddy Tennant
2026-08-30 15:12:52 -04:00
committed by GitHub
parent f5b00de319
commit 1c183e6de9
2 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_ops.py
+4
View File
@@ -978,9 +978,12 @@ class TestOps(unittest.TestCase):
def test_celu(self):
for val in range(1, 5):
helper_test_op([(45,65)], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val))
helper_test_op([(3,3)], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val), low=300, high=400)
helper_test_op([()], lambda x: torch.nn.functional.celu(x,val), lambda x: x.celu(val))
def test_selu(self):
helper_test_op([(45,65)], torch.nn.functional.selu, Tensor.selu)
helper_test_op([(3,3)], torch.nn.functional.selu, Tensor.selu, low=300, high=400)
helper_test_op(None, torch.nn.functional.selu, Tensor.selu, vals=[[-1.,0.,1.]])
helper_test_op([()], torch.nn.functional.selu, Tensor.selu)
def test_silu(self):
helper_test_op([(45,65)], torch.nn.functional.silu, Tensor.silu)
@@ -1117,6 +1120,7 @@ class TestOps(unittest.TestCase):
def test_elu(self):
helper_test_op([(45,65)], torch.nn.functional.elu, Tensor.elu)
helper_test_op([(45,65)], lambda x: torch.nn.functional.elu(x, alpha=0.1), lambda x: Tensor.elu(x, alpha=0.1))
helper_test_op([(3,3)], torch.nn.functional.elu, Tensor.elu, low=300, high=400)
helper_test_op([()], torch.nn.functional.elu, Tensor.elu)
def test_relu6(self):
helper_test_op([(45,65)], torch.nn.functional.relu6, Tensor.relu6)
tinygrad/mixin/elementwise.py
+3 -3
View File
@@ -971,7 +971,7 @@ class ElementwiseMixin(CreationMixin):
print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).elu().numpy())
```
"""
return self.relu() - alpha*(1-self.exp()).relu()
return (self > 0).where(self, alpha*((self - self.relu()).exp() - 1))
def celu(self, alpha=1.0) -> Self:
"""
@@ -983,7 +983,7 @@ class ElementwiseMixin(CreationMixin):
print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).celu().numpy())
```
"""
return self.maximum(0) + (alpha * ((self / alpha).exp() - 1)).minimum(0)
return alpha * (self / alpha).elu()
def selu(self, alpha=1.67326, gamma=1.0507) -> Self:
"""
@@ -995,7 +995,7 @@ class ElementwiseMixin(CreationMixin):
print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).selu().numpy())
```
"""
return gamma * (self >= 0).where(self, alpha * (self.exp() - 1))
return gamma * self.elu(alpha)
def softplus(self, beta=1.0) -> Self:
"""