zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-02 13:54:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: Added "tanh" option to GELU approximation (#1268)

Browse Source
This commit is contained in:
Atakan Tekparmak
2024-07-28 09:07:56 +02:00
committed by GitHub
parent 8cfb9fc0b8
commit 6e06e3a904
2 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/tests/test_nn.py
View File

@@ -722,8 +722,13 @@ class TestLayers(mlx_tests.MLXTestCase):
out = nn.GELU()(mx.array(inputs))
self.assertTrue(np.allclose(out, expected))
# Test the precise/tanh approximation
out_approx = nn.GELU(approx="precise")(mx.array(inputs))
out_approx_tanh = nn.GELU(approx="tanh")(mx.array(inputs))
self.assertTrue(np.allclose(out_approx, expected_approx))
self.assertTrue(np.allclose(out_approx_tanh, expected_approx))
self.assertTrue(np.allclose(out_approx, out_approx_tanh))
# Crudely check the approximations
x = mx.arange(-6.0, 6.0, 12 / 100)