Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix some sigmoid extreme #608

Closed
wants to merge 3 commits into from
Closed

fix some sigmoid extreme #608

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
908aa87
fix some sigmoid extreme
chenyuxyz Dec 14, 2024
0344c68
test_tanh_extreme
chenyuxyz Dec 14, 2024
8f5455d
just sigmoid gradient
chenyuxyz Dec 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 4 additions & 5 deletions test/test_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -807,14 +807,16 @@ def test_softsign_exact(self):
def test_sigmoid(self):
helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid)
helper_test_op([()], torch.sigmoid, Tensor.sigmoid)
@unittest.skip("TODO: fix sigmoid stability")
def test_sigmoid_extreme(self):
helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid, low=300, high=400)
helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid, low=-400, high=-300)
x = Tensor([300.0])
self.assertAlmostEqual(x.sigmoid()[0].gradient(x)[0].item(), 0.0)
x = Tensor([-300.0])
self.assertAlmostEqual(x.sigmoid()[0].gradient(x)[0].item(), 0.0)
def test_hardsigmoid(self):
helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid)
helper_test_op([()], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid)
@unittest.skip("TODO: fix sigmoid stability")
def test_hardsigmoid_extreme(self):
helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid, low=300, high=400)
helper_test_op([(45,65)], torch.sigmoid, Tensor.sigmoid, low=-400, high=-300)
Expand All @@ -835,14 +837,12 @@ def test_erf(self):

def test_gelu(self):
helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), Tensor.gelu)
@unittest.skip("TODO: fix sigmoid stability")
def test_gelu_extreme(self):
helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), Tensor.gelu, low=300, high=400)
helper_test_op([(45,65)], lambda x: torch.nn.functional.gelu(x, approximate="tanh"), Tensor.gelu, low=-400, high=-300)
def test_quick_gelu(self):
helper_test_op([(45,65)], lambda x: x * torch.sigmoid(1.702 * x), Tensor.quick_gelu)
helper_test_op([()], lambda x: x * torch.sigmoid(1.702 * x), Tensor.quick_gelu)
@unittest.skip("TODO: fix sigmoid stability")
def test_quick_gelu_extreme(self):
helper_test_op([(45,65)], lambda x: x * torch.sigmoid(1.702 * x), Tensor.quick_gelu, low=300, high=400)
helper_test_op([(45,65)], lambda x: x * torch.sigmoid(1.702 * x), Tensor.quick_gelu, low=-400, high=-300)
Expand Down Expand Up @@ -1348,7 +1348,6 @@ def test_cosh(self):
helper_test_op([(45,65)], lambda x: x.cosh(), grad_atol=1e-6, low=300, high=303, forward_only=True)
def test_tanh(self):
helper_test_op([(45,65)], lambda x: x.tanh(), grad_atol=1e-6)
@unittest.skip("TODO: fix sigmoid stability")
def test_tanh_extreme(self):
helper_test_op([(45,65)], lambda x: x.tanh(), grad_atol=1e-6, low=-300, high=-297)
helper_test_op([(45,65)], lambda x: x.tanh(), grad_atol=1e-6, low=300, high=303)
Expand Down
8 changes: 8 additions & 0 deletions tinygrad/codegen/uopgraph.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,10 @@ def threefry2x32(x: UOp, key: UOp):

return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)

# ***** other math rewrite ****

def sigmoid_like(x:UOp, y:UOp): return (t:=(1/(x+1))) * (1-t) * y

# ***** main rewriter *****

def loop_collapse(compval, multconst, rng:UOp, acc:UOp, idx2=None,idx3=None,extra=None,vec=None,ne=None,
Expand Down Expand Up @@ -308,6 +312,10 @@ def reduce_collapse(acc:UOp, ret:UOp, alu:UOp):
(UPat(Ops.SINK, name="root"),
lambda root: UOp(Ops.SINK, root.dtype, tuple(flatten(x.src if x.op in {Ops.SINK, Ops.UNROLL} else (x,) for x in root.src)), root.arg)
if any(x.op in {Ops.SINK, Ops.UNROLL} for x in root.src) else None),
# stable sigmoid
(UPat.var("x")*(((UPat.var("x")+1)*(UPat.var("x")+1)).reciprocal()), lambda x: sigmoid_like(x, x.const_like(1))),
(UPat.var("x")*(((UPat.var("x")+1)*(UPat.var("x")+1)).reciprocal()*UPat.var("y")), sigmoid_like),
(UPat.var("x")*(((UPat.var("x")+1)*(UPat.var("x")+1)*(UPat.var("x")+1)).reciprocal()), lambda x: sigmoid_like(x, (x+1).reciprocal())),
])

# *** uop expander ***
Expand Down
Loading