From aa1a6f2132ef878e7409eecb62e512834d7700fb Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 7 Aug 2025 10:43:18 -0700 Subject: [PATCH] support threshold in Tensor.softplus (#11564) fix gradient for large input --- test/test_ops.py | 5 +++-- tinygrad/tensor.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 8ec9e54885..fdd81c7b9e 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -957,8 +957,9 @@ class TestOps(unittest.TestCase): helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6) helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6) helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=1/3), lambda t: Tensor.softplus(t, beta=1/3), grad_atol=1e-6) - # # TODO: support threshold and enable this - # helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400) + helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3, threshold=0.5), + lambda t: Tensor.softplus(t, beta=3, threshold=0.5), grad_atol=1e-6) + helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=300, high=400) helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6, low=-400, high=-300) helper_test_op([()], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 9bac472c83..f66102bc16 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -3507,15 +3507,16 @@ class Tensor(MathTrait): """ return self * self.softplus().tanh() - def softplus(self, beta=1) -> Tensor: + def softplus(self, beta=1.0, threshold=20.0) -> Tensor: """ Applies the Softplus function element-wise. + For numerical stability, the implementation folds into identity function when `self * beta > threshold`. ```python exec="true" source="above" session="tensor" result="python" print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).softplus().numpy()) ``` """ - return (1/beta) * (1 + (self*beta).exp()).log() + return (self * beta > threshold).where(self, (1/beta) * (1 + (self*beta).exp()).log()) def softsign(self) -> Tensor: """