From 1c183e6de91ce83affc36f393641baee58f0e78e Mon Sep 17 00:00:00 2001
From: Teddy Tennant <teddytennant@icloud.com>
Date: Sun, 30 Aug 2026 15:12:52 -0400
Subject: [PATCH] fix elu/celu/selu gradient being nan on large inputs (#17852)

* 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
---
 test/backend/test_ops.py      | 4 ++++
 tinygrad/mixin/elementwise.py | 6 +++---
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/test/backend/test_ops.py b/test/backend/test_ops.py
index 1109dbec4b..a47c90ae1d 100644
--- a/test/backend/test_ops.py
+++ b/test/backend/test_ops.py
@@ -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)
diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py
index 326cc3d64d..914b207568 100644
--- a/tinygrad/mixin/elementwise.py
+++ b/tinygrad/mixin/elementwise.py
@@ -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:
     """