zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-06 16:51:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

segfaut layer norm grad (#955)

Browse Source
This commit is contained in:
Awni Hannun
2024-04-04 10:59:15 -07:00
committed by GitHub
parent e142aaf8a1
commit d88d2124b5
2 changed files with 12 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/tests/test_fast.py
View File

@@ -375,6 +375,17 @@ class TestFast(mlx_tests.MLXTestCase):
self.assertLess(mx.abs(gb1).max(), 1e-9)
self.assertLess(mx.abs(gb2).max(), 1e-9)
def test_layer_norm_grad_no_params(self):
eps = 1e-5
f1 = lambda x: layer_norm(x, None, None, eps).sum()
f2 = lambda x: mx.fast.layer_norm(x, None, None, eps).sum()
x = mx.random.normal(shape=(2, 2, 8))
mx.eval(x)
gx1 = mx.grad(f1)(x)
gx2 = mx.grad(f2)(x)
self.assertTrue(mx.allclose(gx1, gx2, atol=1e-6))
def test_layer_norm_grad_params(self):
eps = 1e-5
f1 = lambda params, x: (layer_norm(x, params[0], params[1], eps)).sum()