From 572e0a4ac3365ca1b8fcc208cf5257937f11e6d1 Mon Sep 17 00:00:00 2001
From: Dan Anderson <dan@kd2ycu.com>
Date: Tue, 10 Mar 2026 00:48:45 -0400
Subject: [PATCH] Validate dims in rope (#3230)

Co-authored-by: KD2YCU <me@kd2ycu.com>
---
 mlx/fast.cpp              | 17 +++++++++++++++++
 python/tests/test_fast.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/mlx/fast.cpp b/mlx/fast.cpp
index cd316d49..a668fe9a 100644
--- a/mlx/fast.cpp
+++ b/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;
diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py
index 9d5cfba4..a9b708fe 100644
--- a/python/tests/test_fast.py
+++ b/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)