[Feature] Metal implementation of SVD

[vgoklani]

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.

[awni]

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)

[awni]

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.

[vgoklani]

Thanks!

[ArkAung]

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.

[awni]

I don't know of anyone working on this. Happy to accept a contribution.

[huangrzh]

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?