ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
17.66k stars 1.02k forks source link

[Feature] Metal implementation of SVD #1392

Open vgoklani opened 2 months ago

vgoklani commented 2 months ago

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 commented 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)
awni commented 2 months ago

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 commented 2 months ago

Thanks!

ArkAung commented 1 month ago

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 commented 1 month ago

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