ml-explore
Support for matrix multiplication (matmul/einsum) with complex matrices is missing
Error
ValueError: [matmul] Only real floating point types are supported but complex64 and complex64 were provided which results in complex64, which is not a real floating point type.
Description:
I am using MLX for simulating laser beams and light fields, where the core data structure is a complex matrix. The lack of support for complex matrix multiplication (matmul) is hindering my ability to perform essential operations, particularly in the context of Quantum and Light Physics simulations.
The absence of this functionality is a significant disadvantage, especially because operations like einsum depend on matmul. This limitation severely impacts the ability to perform certain linear algebra operations required for complex physics simulations.
Problem:
In quantum mechanics and light physics, simulations often require matrix operations on complex matrices in multidimensional bases. Currently, MLX lacks support for matmul with complex matrices, which is a fundamental operation for many advanced simulations such as:
- Quantum simulations: Many algorithms in quantum physics rely on matrix multiplication, such as in quantum state propagation and quantum gate operations.
- Laser beam simulations: In simulations involving laser beams, Fourier transforms and spatial-frequency domain manipulations are critical. These operations often require complex matrix multiplication.
- Multidimensional basis operations: Simulating complex light fields or multi-mode systems requires working with multidimensional arrays, where matrix multiplication is crucial for transforming between different bases.
Why This Is Important:
- Laser and light field simulations depend on operations like matrix multiplication to model diffraction, interference, and propagation.
- Quantum physics simulations need complex matrix operations to evolve quantum states or compute matrix elements for Hamiltonians.
- For macOS users MLX would be a preferred library due to its performance benefits, but the lack of matmul for complex matrices makes it challenging to use for these fields.
Suggested Solution:
It would be great if MLX could implement support for matrix multiplication on complex matrices, particularly for multidimensional arrays. This would enable advanced simulations in physics, like laser propagation and quantum mechanics, to be performed with the full power of MLX's optimized performance.
Possible Workaround:
If there is any workaround or alternative way to perform matrix multiplication on complex matrices in MLX (or any other MLX-compatible library), I would appreciate any advice or suggestions.
Environment:
- OS: macOS 15.3.2, Apple M4
- MLX Version: 023.2
- Python Version: 3.13
Thank you for your work on MLX, and I look forward to hearing if this issue could be addressed in a future release.
Related issues: [Feature request] Improved support for complex64 arrays #361 Add Gemm for complex64 #140
It should be quite doable to add a complex GEMM for the CPU backend using BLAS (zgemm or something)
As for the Metal backend it will probably be more work but also doable. I'm not sure if it makes sense to go through steel for this or write a custom kernel (CC @jagrit06 might have something to add there).
The Metal backend definitely seems more appealing to me, as my reason for switching was to move from NumPy to MLX and take advantage of greater performance by utilizing the GPU as well. Similar alternatives are used by owners of CUDA-compatible GPUs for simulations with CuPy.