default install cuda on linux (#2462)

This commit is contained in:
Awni Hannun 2025-08-04 15:33:05 -07:00 committed by GitHub
parent 737dd6d1ac
commit 6ad0889c8a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 12 additions and 15 deletions

View File

@ -75,16 +75,11 @@ macOS, run:
pip install mlx
```
To install the CUDA backend on Linux, run:
On Linux, the same command (``pip install mlx``) will install the CUDA backend
by default. To install a CPU-only Linux package, run:
```bash
pip install "mlx[cuda]"
```
To install a CPU-only Linux package, run:
```bash
pip install "mlx[cpu]"
pip install mlx[cpu]
```
Checkout the

View File

@ -30,7 +30,7 @@ MLX has a CUDA backend which you can install with:
.. code-block:: shell
pip install "mlx[cuda]"
pip install mlx
To install the CUDA package from PyPi your system must meet the following
requirements:
@ -49,7 +49,7 @@ For a CPU-only version of MLX that runs on Linux use:
.. code-block:: shell
pip install "mlx[cpu]"
pip install mlx[cpu]
To install the CPU-only package from PyPi your system must meet the following
requirements:

View File

@ -274,11 +274,13 @@ if __name__ == "__main__":
# - Package name is back-end specific, e.g mlx-metal
if build_stage != 2:
if build_stage == 1:
install_requires.append(
f'mlx-metal=={version}; platform_system == "Darwin"'
)
extras["cuda"] = [f'mlx-cuda=={version}; platform_system == "Linux"']
extras["cpu"] = [f'mlx-cpu=={version}; platform_system == "Linux"']
install_requires += [
f'mlx-metal=={version}; platform_system == "Darwin"',
f'mlx-cuda=={version}; extra != "cpu" and platform_system == "linux"',
]
extras["cpu"] = [
f'mlx-cpu=={version}; extra == "cpu" and platform_system == "linux"'
]
_setup(
name="mlx",