From 1d9076c24b699f4202471f43e53448e811ade932 Mon Sep 17 00:00:00 2001
From: thesuryash <suryashmalviya@gmail.com>
Date: Fri, 23 May 2025 08:16:33 -0400
Subject: [PATCH] Implementing Complex Matmul using Karatsuba Algorithm

---
 mlx/ops.cpp | 34 ++++++++++++++++++++++------------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index a72c2bc85..1ba6cba07 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -2862,20 +2862,30 @@ array matmul(
         << " second input with shape " << b.shape() << ".";
     throw std::invalid_argument(msg.str());
   }
+
+// complex matmul using Karatsuba's Algorithm
+
+if (a.dtype() == complex64 && b.dtype() == complex64) {
+  // Extract real and imaginary parts
+  auto a_real = real(a, s);
+  auto a_imag = imag(a, s);
+  auto b_real = real(b, s);
+  auto b_imag = imag(b, s);
+
+  // Compute real and imaginary components of the result
+  auto m1 = matmul(a_real, b_real, s);
+  auto m2 = matmul(a_imag, b_imag, s);
+  auto m3 = matmul(add(a_real, a_imag, s), add(b_real, b_imag, s), s);
+
+  auto c_real = subtract(m1, m2, s);
+  auto c_imag = subtract(m3, add(m1, m2, s), s);
+
+  return add(c_real, multiply(array(complex64_t{0, 1}, complex64), c_imag, s), s);
+}
+
   // Type promotion
   auto out_type = promote_types(a.dtype(), b.dtype());
-  // Complex matmul in terms of real matmuls
-  if (out_type == complex64) {
-    auto a_real = real(a, s);
-    auto b_real = real(b, s);
-    auto a_imag = imag(a, s);
-    auto b_imag = imag(b, s);
-    auto c_real =
-        subtract(matmul(a_real, b_real, s), matmul(a_imag, b_imag, s), s);
-    auto c_imag = add(matmul(a_real, b_imag, s), matmul(a_imag, b_real, s), s);
-    return add(
-        c_real, multiply(array(complex64_t{0, 1}, complex64), c_imag, s), s);
-  }
+ 
 
   if (!issubdtype(out_type, floating)) {
     std::ostringstream msg;