error in rms for wrong size (#1562)

2025-10-31 16:21:27 +08:00 · 2024-11-04 13:24:02 -08:00
parent f1951d6cce
commit 76f275b4df
2 changed files with 13 additions and 0 deletions
--- a/python/tests/test_fast.py
+++ b/python/tests/test_fast.py
@@ -308,6 +308,11 @@ class TestFast(mlx_tests.MLXTestCase):
        rx_fast = mx.fast.rms_norm(x, weight, eps)
        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-6)

+        # Wrong size w raises
+        with self.assertRaises(ValueError):
+            x = mx.random.uniform(shape=(1, 5))
+            mx.fast.rms_norm(x, mx.ones((4,)), 1e-5)
+
    def test_rms_norm_grad(self):
        D = 32
        eps = 1e-5