Simple way to install mlx[cuda] with uv?

[arogozhnikov]

Hello, I am testing (again) MLX with CUDA.

I found installation (I use uv, maybe that matters) a bit finicky - in some combination, if mlx is installed after torch, it would work; sometimes it doesn't. Some cases are totally weird:

uv run --with 'mlx[cuda]==0.30.0' python -c 'import mlx.core as mx;  c = mx.random.normal([32, 32], stream=mx.gpu); mx.eval(c)'
# above works. Using `c@c` also works while I'd expect this operation to run on gpu. 

let's change command to

# test_mlx.py
import mlx.core as mx

print(mx.default_device())  # confirm gpu
x = mx.random.normal([1024, 1024])
mx.eval(mx.addmm(x[0], x, x))

uv run --with 'mlx[cuda]==0.30.0' python test_mlx.py

adding addmm makes code fail with

RuntimeError: Environment variable CUDA_HOME or CUDA_PATH is not set.

Ok, let's add toolkit with all packages, to match torch installation:

uv run --with 'mlx[cuda]==0.30.0' --with 'cuda-toolkit[all]==12.9' --with 'nvidia-cudnn-cu12' python test_mlx.py
# same: CUDA_HOME or CUDA_PATH is not set.
uv run --with 'mlx[cuda]==0.29.0' --with 'cuda-toolkit[all]==12.9' --with 'nvidia-cudnn-cu12' python test_mlx.py
# fails with ImportError: libnccl.so: cannot open shared object file: No such file or directory
# python nccl package is already installed.

I was able to make it work with one of previous MLX versions 😅 .

I'd appreciate help with getting right installation which avoids installing system packages, as I don't want to keep multiple containers while running simple tests.

[zcbenz]

Currently mlx-cuda does require a system installed CUDA toolkit, your use case is definitely reasonable and I have opened an issue for it: #2842.