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()