zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Layer norm throws on dimension mismatch (#2870)

Browse Source
This commit is contained in:
Awni Hannun
2025-12-04 11:21:05 -08:00
committed by GitHub
parent 9abb0b8123
commit 941cfe23d7
2 changed files with 42 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
python/tests/test_fast.py
View File

@@ -103,9 +103,6 @@ class TestFast(mlx_tests.MLXTestCase):
dims, _, base, scale, offset, _ = defaults
for dtype in dtypes:
x = mx.random.uniform(shape=(2, T, dims)).astype(dtype)
ry = rope_orig(
x.astype(mx.float32), dims, traditional, base, scale, offset
)
rx = rope_orig(x, dims, traditional, base, scale, offset)
rx_fast = mx.fast.rope(
x,
@@ -116,9 +113,10 @@ class TestFast(mlx_tests.MLXTestCase):
offset=offset,
)
if dtype != mx.float32:
self.assertLessEqual(
mx.abs(ry - rx_fast).max(), mx.abs(ry - rx).max()
ry = rope_orig(
x.astype(mx.float32), dims, traditional, base, scale, offset
)
self.assertLess(mx.abs(ry - rx_fast).max(), tolerances[dtype])
self.assertLess(mx.abs(rx - rx_fast).max(), tolerances[dtype])
dims, dtype, base, scale, _, _ = defaults
@@ -455,6 +453,17 @@ class TestFast(mlx_tests.MLXTestCase):
self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5)
self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5)
def test_layer_norm_dim_check(self):
with self.assertRaises(ValueError):
weight = mx.ones((129,))
x = mx.random.randint(low=0, high=10, shape=(4, 128))
mx.fast.layer_norm(x, weight, None, 1e-3)
with self.assertRaises(ValueError):
bias = mx.ones((129,))
x = mx.random.randint(low=0, high=10, shape=(4, 128))
mx.fast.layer_norm(x, None, bias, 1e-3)
def test_layer_norm(self):
# Per dtype absolute tolerance
tolerances = {mx.float32: 1e-5, mx.float16: 5e-3, mx.bfloat16: 5e-2}