From bf6421bfb05d9e39d182432611b99ee9a3dca933 Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Fri, 8 May 2026 00:25:27 -0700
Subject: [PATCH] Add test

---
 python/tests/test_fast.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/python/tests/test_fast.py b/python/tests/test_fast.py
index a9b708fe..0c705733 100644
--- a/python/tests/test_fast.py
+++ b/python/tests/test_fast.py
@@ -365,6 +365,28 @@ class TestFast(mlx_tests.MLXTestCase):
         rx = rope_orig(x, dims, traditional, base, scale, offset)
         self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
 
+    def test_rope_single_batch(self):
+        base = 10000.0
+        scale = 1.0
+        offset = 5
+
+        for traditional in [True, False]:
+            for B in [2, 4, 8]:
+                for n_head in [1, 4, 7]:
+                    for dims in [64, 128]:
+                        x = mx.random.uniform(shape=(B, n_head, 1, dims))
+                        mx.eval(x)
+                        rx_fast = mx.fast.rope(
+                            x,
+                            dims,
+                            traditional=traditional,
+                            base=base,
+                            scale=scale,
+                            offset=offset,
+                        )
+                        rx = rope_orig(x, dims, traditional, base, scale, offset)
+                        self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5)
+
     def test_rope_with_large_offset(self):
         x = mx.random.normal(shape=(1, 1, 1024, 32))
         rx_fp32 = mx.fast.rope(