Apple MLX Contribution
Merged PR #2220 — faster complex matrix multiplication • 2025
I got production code merged into a major open-source ML framework maintained by Apple: PR #2220 makes MLX's complex matrix multiplication measurably faster on both CPU and GPU.
Problem
MLX computed a complex matrix product (A + Bi)(C + Di) the textbook way: four real matrix multiplications (AC, BD, AD, BC) plus additions. Matrix multiplication is the expensive part — every multiply you avoid is real time saved on every complex matmul in the framework.
What I built
A Karatsuba-style reformulation. Compute three products — AC, BD, and (A + B)(C + D) — then recover the real part as AC − BD and the imaginary part as (A + B)(C + D) − AC − BD. Three matrix multiplications instead of four, at the cost of a few extra matrix additions.
The trade wins because multiplication dominates: for n×n matrices a matmul costs O(n³) while the added additions cost O(n²), so the saved multiplication pays for the bookkeeping and the gap widens as matrices grow. I benchmarked the change with MLX's own benchmarking utilities and worked through Apple's review process — debugging, tests, and code review — until it merged.
Outcome
- Merged into Apple's official MLX repository as PR #2220
- Complex matrix multiplication now uses 3 internal multiplications instead of 4
- Benchmarks reported in the PR notes: ~17.70% CPU speedup and ~19.30% GPU speedup for the tested configuration