ml-explore
[Feature] Metal implementation of SVD
Describe the bug The mx.linalg.svd crashes
To Reproduce
Include code snippet
import mlx.core as mx
a = mx.array([ [1., 2], [2, 3] ])
u, singular_values, v_transpose = mx.linalg.svd(a=a)
libc++abi: terminating due to uncaught exception of type std::runtime_error: [SVD::eval_gpu] Metal SVD NYI.
Expected behavior
import numpy as np
import mlx.core as mx
a = mx.array([ [1., 2], [2, 3] ])
u, singular_values, v = np.linalg.svd(a=a)
singular_values array([4.236068 , 0.23606798], dtype=float32)
Desktop (please complete the following information): MLX 0.17.2
Additional context Add any other context about the problem here.
That is unfortunately expected behavior. Right now svd (and several other linalg operations are only supported on the CPU back-end. You can fix that by passing in the CPU stream to the operation like so:
import mlx.core as mx
a = mx.array([ [1., 2], [2, 3] ])
u, singular_values, v_transpose = mx.linalg.svd(a=a, stream=mx.cpu)
I'm going to change this from a bug to a feature request and mark it as such. Note it's not a trivial op to implement on the GPU so it could take some time to get this feature into MLX.
Is anyone working on this? If not, I would like to start working on this. Happy to collaborate with other contributors who would like to collab on the PR.
Is anyone working on this? If not, I would like to start working on this. Happy to collaborate with other contributors who would like to collab on the PR.
Hi, how is it going with SVD?
Hello, I'd like to ask if there has been any progress on the GPU SVD feature?
I was working on it for a while and got side-tracked because of other responsibilities. If people want this, I will properly dedicate this and revive the draft PR (now stale and closed). There are some challenges with making it work for large matrices. I am willing to take co-contributors to get this across the line.