From d4aafb9161d90c38eb7f06565374ee5a76c66c26 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Mon, 2 Jun 2025 13:43:22 -0700 Subject: [PATCH] fix --- mlx/ops.cpp | 3 +-- python/tests/test_blas.py | 14 +++++++------- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/mlx/ops.cpp b/mlx/ops.cpp index cc1958572..9602f667a 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2864,8 +2864,7 @@ array matmul( } // complex matmul using Karatsuba's Algorithm - - if (a.dtype() == complex64 && b.dtype() == complex64) { + if (a.dtype() == complex64 || b.dtype() == complex64) { // Extract real and imaginary parts auto a_real = real(a, s); auto a_imag = imag(a, s); diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index df459eadc..8c7a97ba8 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -1210,13 +1210,6 @@ class TestBlas(mlx_tests.MLXTestCase): self.assertTrue(np.allclose(c, c_np)) # Test addmm - M = 16 - K = 50 - N = 32 - - def rand(shape): - return mx.random.uniform(shape=shape) + 1j * mx.random.uniform(shape=shape) - a = rand((M, K)) b = rand((K, N)) c = rand((M, N)) @@ -1224,6 +1217,13 @@ class TestBlas(mlx_tests.MLXTestCase): out_np = 2.0 * np.matmul(a, b) + 2.0 * c self.assertTrue(np.allclose(out, out_np)) + # complex with real + a = rand((M, K)).real + b = rand((K, N)) + c = mx.matmul(a, b) + c_np = np.matmul(a, b) + self.assertTrue(np.allclose(out, out_np)) + if __name__ == "__main__": unittest.main()