Custom primitive + RoPE fat op (#676)

* extensions start

* rope custom op

* fix build

* docs + rope benchmark

* fix test

* Add a Metal kernel for RoPE

* Fix position of traditional

* transform tests

* Move rope computation to float and fix tests

* Fix the test and a typo

* change to fast

* fix no metal build

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

benchmarks/python/rope_bench.py

@@ -0,0 +1,35 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+from time_utils import time_fn
+
+
+def time_rope():
+    rope = nn.RoPE(4096)
+
+    # vec
+    x = mx.random.uniform(shape=(1, 4096)).astype(mx.float16)
+    mx.eval(x)
+
+    def rope_vec(x):
+        for _ in range(32):
+            x = rope(x)
+        return x
+
+    time_fn(rope_vec, x)
+
+    # matrix
+    x = mx.random.uniform(shape=(1024, 4096)).astype(mx.float16)
+    mx.eval(x)
+
+    def rope_mat(x):
+        for _ in range(32):
+            x = rope(x)
+        return x
+
+    time_fn(rope_mat, x)
+
+
+if __name__ == "__main__":
+    time_rope()