From e5b24608df07191f059a7aa750a7dca5b77cd285 Mon Sep 17 00:00:00 2001
From: Tobias Buck <buck@mpia.de>
Date: Fri, 16 Jan 2026 10:29:39 +0100
Subject: [PATCH 1/3] Move jax dependency for CPU from optional to default

---
 pyproject.toml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

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",
     ]

From 4d04416b44420fe79a8e8ecc30f0218c63edd297 Mon Sep 17 00:00:00 2001
From: Tobias Buck <buck@mpia.de>
Date: Fri, 16 Jan 2026 10:33:14 +0100
Subject: [PATCH 2/3] Update installation instructions to set CPU as default
 and clarify GPU setup

---
 README.md            | 8 ++++----
 docs/CONTRIBUTING.md | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 77df351..e823213 100644
--- a/README.md
+++ b/README.md
@@ -26,15 +26,15 @@ Key features include:
 
 ## Installation
 
-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:
+The Python package `rubix` is published on GitHub and can be installed alongside its runtime dependencies (including JAX) by choosing the relevant extras. The CPU-only version will be installed by default as follows:
 
 ```
 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
 
@@ -44,7 +44,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
 

From 3a0c7e7fc5b08c0f1e8db7573d162ce87af391fd Mon Sep 17 00:00:00 2001
From: anschaible <131476730+anschaible@users.noreply.github.com>
Date: Fri, 16 Jan 2026 10:53:39 +0100
Subject: [PATCH 3/3] Update README with installation instructions

Added instructions for GPU acceleration and CPU-only installation.
---
 README.md | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/README.md b/README.md
index 6eed122..241a13a 100644
--- a/README.md
+++ b/README.md
@@ -31,6 +31,12 @@ 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:
 
 ```