zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-12 04:06:39 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

default install cuda on linux (#2462)

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
README.md
View File

@ -75,16 +75,11 @@ macOS, run:
pip install mlx 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 ```bash
pip install "mlx[cuda]" pip install mlx[cpu]
```
To install a CPU-only Linux package, run:
```bash
pip install "mlx[cpu]"
``` ```
Checkout the Checkout the

4
docs/src/install.rst
View File

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

12
setup.py
View File

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