revert default cuda install

README.md

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

docs/src/install.rst

@@ -30,7 +30,7 @@ MLX has a CUDA backend which you can install with:
 
 .. code-block:: shell
 
-    pip install mlx
+    pip install "mlx[cuda]"
 
 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:

setup.py

@@ -274,13 +274,11 @@ if __name__ == "__main__":
     #  - Package name is back-end specific, e.g mlx-metal
     if build_stage != 2:
         if build_stage == 1:
-            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"'
-            ]
+            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"']
 
         _setup(
             name="mlx",