diff --git a/README.md b/README.md
index d9796e1..241a13a 100644
--- a/README.md
+++ b/README.md
@@ -31,15 +31,21 @@ Rubix is on [PyPi](https://pypi.org/project/astro-rubix/). You can install it vi
 pip install astro-rubix
 ```
 
+If you need GPU acceleration, you can install it via
+
+```
+pip install astro-rubix[cuda]
+```
+
 Alternatively, the Python package `rubix` is published on GitHub and can be installed alongside its runtime dependencies (including JAX) by choosing the relevant extras. For a CPU-only environment, install with:
 
 ```
 git clone https://github.com/AstroAI-Lab/rubix.git
 cd rubix
-pip install .[cpu]
+pip install .
 ```
 
-If you need GPU acceleration, replace `[cpu]` with `[cuda]` (or install `jax[cuda]` following the [JAX instructions](https://github.com/google/jax#installation) before installing Rubix). The plain `pip install .` command installs the minimal package without JAX and will raise `ImportError` if you try to import `rubix` before adding `jax` manually.
+If you need GPU acceleration, please add the optional dependence with `[cuda]` (or install `jax[cuda]` following the [JAX instructions](https://github.com/google/jax#installation) before installing Rubix).
 
 ## Development installation
 
@@ -49,7 +55,7 @@ the following editable installation from this repository:
 ```
 git clone https://github.com/AstroAI-Lab/rubix.git
 cd rubix
-python -m pip install --editable .[cpu,tests,dev]
+python -m pip install --editable .[tests,dev]
 ```
 
 Having done so, the test suite can be run using `pytest`:
diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md
index 0863887..85403b8 100644
--- a/docs/CONTRIBUTING.md
+++ b/docs/CONTRIBUTING.md
@@ -30,10 +30,10 @@ You can then clone the repo and install it in editable mode with the extra devel
 ```bash
 git clone https://github.com/AstroAI-Lab/rubix
 cd rubix
-pip install -e .[cpu,dev]
+pip install -e .[dev]
 ```
 
-Note: if you are planning to use RUBIX on the GPU you need to replace the `cpu` dependency with the corresponding `cuda` dependency. If you additionally plan to build the docs locally you'll also need to include the `docs` dependency group.
+Note: if you are planning to use RUBIX on the GPU you need to add the corresponding optional `cuda` dependency. If you additionally plan to build the docs locally you'll also need to include the `docs` dependency group.
 
 ### Setting up pre-commit hooks
 
diff --git a/pyproject.toml b/pyproject.toml
index a7c4449..8184f57 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,6 +59,7 @@ keywords =[
 urls = { Homepage = "https://astro-rubix.web.app", Repository = "https://github.com/AstroAI-Lab/rubix", Issues = "https://github.com/AstroAI-Lab/rubix/issues" }
 
 dependencies = [
+    "jax[cpu]>0.5.1",
     "requests",
     "h5py",
     "astropy",
@@ -102,9 +103,7 @@ docs = [
     "sphinx_mdinclude",
     "sphinx_rtd_theme",
 ]
-cpu = [
-    "jax[cpu]>0.5.1",
-    ]
+
 cuda = [
     "jax[cuda]>0.5.1",
     ]