Validate dims in rope (#3230)

Co-authored-by: KD2YCU <[email protected]>

mlx/fast.cpp

@@ -406,6 +406,23 @@ array rope(
   if (offset.dtype().size() != 4) {
     inputs[1] = astype(offset, int32, s);
   }
+  if (dims <= 0) {
+    std::ostringstream msg;
+    msg << "[rope] dims must be positive but got " << dims << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (dims % 2 != 0) {
+    std::ostringstream msg;
+    msg << "[rope] dims must be even but got " << dims << ".";
+    throw std::invalid_argument(msg.str());
+  }
+  if (dims > x.shape(-1)) {
+    std::ostringstream msg;
+    msg << "[rope] dims must not exceed the input's last dimension ("
+        << x.shape(-1) << ") but got " << dims << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
   if (inputs.size() == 3 &&
       (inputs[2].ndim() != 1 || inputs[2].shape(0) != dims / 2)) {
     std::ostringstream msg;

python/tests/test_fast.py

@@ -169,6 +169,41 @@ class TestFast(mlx_tests.MLXTestCase):
                 x, dims, traditional=traditional, base=base, scale=scale, offset=offset
             )
 
+    def test_rope_dims_validation(self):
+        T = 4
+        feature_dim = 64
+        x = mx.random.uniform(shape=(1, T, feature_dim))
+
+        # dims = 0 should raise
+        with self.assertRaises(ValueError):
+            mx.fast.rope(
+                x, dims=0, traditional=False, base=10000.0, scale=1.0, offset=0
+            )
+
+        # negative dims should raise
+        with self.assertRaises(ValueError):
+            mx.fast.rope(
+                x, dims=-2, traditional=False, base=10000.0, scale=1.0, offset=0
+            )
+
+        # odd dims should raise
+        with self.assertRaises(ValueError):
+            mx.fast.rope(
+                x, dims=7, traditional=False, base=10000.0, scale=1.0, offset=0
+            )
+
+        # dims > feature_dim should raise
+        with self.assertRaises(ValueError):
+            mx.fast.rope(
+                x, dims=128, traditional=False, base=10000.0, scale=1.0, offset=0
+            )
+
+        # valid dims should not raise
+        mx.fast.rope(x, dims=32, traditional=False, base=10000.0, scale=1.0, offset=0)
+        mx.fast.rope(
+            x, dims=feature_dim, traditional=False, base=10000.0, scale=1.0, offset=0
+        )
+
     def test_rope_with_freqs(self):
         mx.random.seed(0)