Open vgoklani opened 2 months ago
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.
Thanks!
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.
I don't know of anyone working on this. Happy to accept a contribution.
Describe the bug The mx.linalg.svd crashes
To Reproduce
Include code snippet
Expected behavior
Desktop (please complete the following information): MLX 0.17.2
Additional context Add any other context about the problem here.